From b04afd29ab751e729864bc618dd7b2bf5b270ad7 Mon Sep 17 00:00:00 2001 From: Somshubra Majumdar Date: Fri, 24 May 2024 17:09:18 -0700 Subject: [PATCH] Guard cuda memory allocator update (#9312) * Guard cuda memory allocator update Signed-off-by: smajumdar * Apply isort and black reformatting Signed-off-by: titu1994 --------- Signed-off-by: smajumdar Signed-off-by: titu1994 --- nemo/collections/common/data/lhotse/dataloader.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/nemo/collections/common/data/lhotse/dataloader.py b/nemo/collections/common/data/lhotse/dataloader.py index 32bbc1f3e8f4..01bf51b0e2c6 100644 --- 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: