fix batch_size insistent in report (pytorch#7820)

ManfeiBai · Aug 12, 2024 · f0bbaeb · f0bbaeb
1 parent 8f79488
commit f0bbaeb
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/benchmarks/experiment_runner.py b/benchmarks/experiment_runner.py
@@ -286,6 +286,10 @@ def run_single_config(self):
       reset_rng_state(experiment)
       model = self.model_loader.load_model(model_config, experiment)
 
+      # real batch_size can be updated after load_model, need to update
+      # so the config can be reflected in the report.
+      experiment_config['batch_size'] = experiment.batch_size
+
       # Repeat the experiment and accumulate metrics.
       with model.pick_grad():
         for repeat_iteration in range(self._args.repeat):