Skip to content

Commit

Permalink
reduce max split size
Browse files Browse the repository at this point in the history
  • Loading branch information
priyakasimbeg committed Oct 7, 2023
1 parent 416b88d commit 4600d78
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion submission_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -228,7 +228,7 @@ def train_once(
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:512'
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:256'
logging.info('Initializing optimizer.')
with profiler.profile('Initializing optimizer'):
optimizer_state = init_optimizer_state(workload,
Expand Down

0 comments on commit 4600d78

Please sign in to comment.