minor

mlcommons · Jul 27, 2023 · b047468 · b047468
1 parent ec147ed
commit b047468
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/submission_runner.py b/submission_runner.py
@@ -214,7 +214,7 @@ def train_once(
       #   logging.info('Performing `torch.compile`.')
       #   model_params = torch.compile(model_params)
 
-      model_params = torch.compile(model_params, fullgraph=False, backend='eager')
+      model_params = torch.compile(model_params, fullgraph=False, backend='aot_autograd')
 
   logging.info('Initializing optimizer.')
   with profiler.profile('Initializing optimizer'):