Guard cuda memory allocator update (#9312) (#9313)

* Guard cuda memory allocator update * Apply isort and black reformatting --------- Signed-off-by: smajumdar <[email protected]> Signed-off-by: titu1994 <[email protected]> Co-authored-by: Somshubra Majumdar <[email protected]> Signed-off-by: Jan Lasek <[email protected]>
NVIDIA · Jun 12, 2024 · 234b53f · 234b53f
1 parent a5add76
commit 234b53f
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/nemo/collections/common/data/lhotse/dataloader.py b/nemo/collections/common/data/lhotse/dataloader.py
@@ -531,7 +531,13 @@ def maybe_set_cuda_expandable_segments(enabled: bool):
             warnings.warn(
                 "You have set PYTORCH_CUDA_ALLOC_CONF without expandable_segments:True option. We're setting that option anyway. To disable it, set cuda_expandable_segments=False in NeMo dataloader configuration."
             )
-        torch.cuda.memory._set_allocator_settings("expandable_segments:True")
+
+        try:
+            torch.cuda.memory._set_allocator_settings("expandable_segments:True")
+        except RuntimeError:
+            logging.info(
+                "Failed to set expandable_segments:True for PyTorch CUDA allocator. You may get training speed improvements if you enable this"
+            )
 
 
 def _select_channel(cut, channel_selector: int | str) -> list: