From 2d7ab065030462f151f09ef91f86d3f0f4e6bc62 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Thu, 9 Jun 2022 14:30:21 -0700 Subject: [PATCH] Move aot-autograd opt after model metadata used to setup data config in benchmark.py --- benchmark.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/benchmark.py b/benchmark.py index e2370dcc..f348fcb9 100755 --- a/benchmark.py +++ b/benchmark.py @@ -229,14 +229,14 @@ class BenchmarkRunner: if torchscript: self.model = torch.jit.script(self.model) self.scripted = True - if aot_autograd: - assert has_functorch, "functorch is needed for --aot-autograd" - self.model = memory_efficient_fusion(self.model) - data_config = resolve_data_config(kwargs, model=self.model, use_test_size=not use_train_size) self.input_size = data_config['input_size'] self.batch_size = kwargs.pop('batch_size', 256) + if aot_autograd: + assert has_functorch, "functorch is needed for --aot-autograd" + self.model = memory_efficient_fusion(self.model) + self.example_inputs = None self.num_warm_iter = num_warm_iter self.num_bench_iter = num_bench_iter