fix

mosaicml · Mar 29, 2024 · ce2dff6 · ce2dff6
1 parent ebcdc96
commit ce2dff6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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