fix

mlcommons · Oct 7, 2023 · 416b88d · 416b88d
1 parent da89a8b
commit 416b88d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/submission_runner.py b/submission_runner.py
@@ -227,8 +227,8 @@ def train_once(
         logging.info('Performing `torch.compile`.')
         model_params = torch.compile(model_params)
     # Temporary fix for Conformer OOM
-    if flags.framework == 'pytorch' and flags.workload == 'librispeech_conformer':
-      os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:256'
+    if FLAGS.framework == 'pytorch' and FLAGS.workload == 'librispeech_conformer':
+      os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
   logging.info('Initializing optimizer.')
   with profiler.profile('Initializing optimizer'):
     optimizer_state = init_optimizer_state(workload,