From 416b88dd2d621b2f99f69d58bde3f0eae216146d Mon Sep 17 00:00:00 2001 From: Priya Kasimbeg Date: Sat, 7 Oct 2023 00:47:45 +0000 Subject: [PATCH] fix --- submission_runner.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/submission_runner.py b/submission_runner.py index 069d2ffda..1460f5573 100644 --- 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,