From 8aa2df2bb14d53f30a922b23f072d7e2eec7871a Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Fri, 31 May 2024 21:37:37 -0700 Subject: [PATCH] Guard cuda memory allocator update (#9312) (#9313) * Guard cuda memory allocator update * Apply isort and black reformatting --------- Signed-off-by: smajumdar Signed-off-by: titu1994 Co-authored-by: Somshubra Majumdar --- 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 32bbc1f3e8f4b..01bf51b0e2c63 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: