Skip to content

Commit

Permalink
Guard cuda memory allocator update (NVIDIA#9312) (NVIDIA#9313)
Browse files Browse the repository at this point in the history
* 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: Boxiang Wang <[email protected]>
  • Loading branch information
2 people authored and BoxiangW committed Jun 5, 2024
1 parent 11c6516 commit 113d220
Showing 1 changed file with 7 additions and 1 deletion.
8 changes: 7 additions & 1 deletion nemo/collections/common/data/lhotse/dataloader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down

0 comments on commit 113d220

Please sign in to comment.