Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

torch.cuda.memory_allocated to return {} if not initialized #51179

Closed
wants to merge 5 commits into from
Closed

torch.cuda.memory_allocated to return {} if not initialized #51179

Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
d9666fb
torch.cuda.memory_allocated to return `{}` if not initialized
malfet Jan 27, 2021
7e0d818
Add handling for torch.cuda.device as an argument
malfet Jan 27, 2021
8f3a3b0
Added test that verifies torch.cuda.memory_allocated works with torch…
malfet Jan 27, 2021
63bc8ec
Update test
malfet Jan 28, 2021
776ed5c
Move the test to TestCuda vs TestCudaComm class
malfet Jan 28, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
10 changes: 6 additions & 4 deletions torch/cuda/memory.py
View file
Open in desktop
Expand Up @@ -193,6 +193,8 @@ def _recurse_add_to_result(prefix, obj):

def memory_stats_as_nested_dict(device: Union[Device, int] = None) -> Dict[str, Any]:
r"""Returns the result of :func:`~torch.cuda.memory_stats` as a nested dictionary."""
if not is_initialized():
return {}
device = _get_device_index(device, optional=True)
return torch._C._cuda_memoryStats(device)

Expand Down Expand Up @@ -303,7 +305,7 @@ def memory_allocated(device: Union[Device, int] = None) -> int:
needs to be created on GPU. See :ref:`cuda-memory-management` for more
details about GPU memory management.
"""
return memory_stats(device=device)["allocated_bytes.all.current"]
return memory_stats(device=device).get("allocated_bytes.all.current", 0)


def max_memory_allocated(device: Union[Device, int] = None) -> int:
Expand All @@ -325,7 +327,7 @@ def max_memory_allocated(device: Union[Device, int] = None) -> int:
See :ref:`cuda-memory-management` for more details about GPU memory
management.
"""
return memory_stats(device=device)["allocated_bytes.all.peak"]
return memory_stats(device=device).get("allocated_bytes.all.peak", 0)


def memory_reserved(device: Union[Device, int] = None) -> int:
Expand All @@ -341,7 +343,7 @@ def memory_reserved(device: Union[Device, int] = None) -> int:
See :ref:`cuda-memory-management` for more details about GPU memory
management.
"""
return memory_stats(device=device)["reserved_bytes.all.current"]
return memory_stats(device=device).get("reserved_bytes.all.current", 0)


def max_memory_reserved(device: Union[Device, int] = None) -> int:
Expand All @@ -363,7 +365,7 @@ def max_memory_reserved(device: Union[Device, int] = None) -> int:
See :ref:`cuda-memory-management` for more details about GPU memory
management.
"""
return memory_stats(device=device)["reserved_bytes.all.peak"]
return memory_stats(device=device).get("reserved_bytes.all.peak", 0)


def memory_cached(device: Union[Device, int] = None) -> int:
Expand Down