From ce2dff68495d65acb09cea50c9e1d07f0d198a8f Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 29 Mar 2024 00:33:34 -0700 Subject: [PATCH] fix --- scripts/train/train.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/scripts/train/train.py b/scripts/train/train.py index 01a4f78c24..615be2cb45 100644 --- a/scripts/train/train.py +++ b/scripts/train/train.py @@ -141,7 +141,8 @@ def main(cfg: DictConfig) -> Trainer: if expandeable_segments: cuda_alloc_conf.append('expandeable_segments:True') - os.environ['PYTORCH_CUDA_ALLOC_CONF'] = ','.join(cuda_alloc_conf) + if len(cuda_alloc_conf) > 0: + os.environ['PYTORCH_CUDA_ALLOC_CONF'] = ','.join(cuda_alloc_conf) # Set CUDA lazy loading # This can save a bit of memory if not all modules are needed