fixed return type for cuda.memory.mem_get_info() (#81073)
Return type was `int` but function actually returns a tuple of two ints. The first being the free gpu memory in bytes and the second being the total available gpu memory in bytes.
Return type was fixed to correctly read `Tuple[int, int]` and the `Tuple` class was imported from `typing`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/81073
Approved by: https://github.com/ngimel
diff --git a/torch/cuda/memory.py b/torch/cuda/memory.py
index a4b7b1d..904ad30 100644
--- a/torch/cuda/memory.py
+++ b/torch/cuda/memory.py
@@ -1,7 +1,7 @@
import collections
import contextlib
import warnings
-from typing import Any, Dict, Union
+from typing import Any, Dict, Union, Tuple
import torch
from . import is_initialized, _get_device_index, _lazy_init
@@ -573,7 +573,7 @@
lines.append(f"process {p.pid:>10d} uses {mem:>12.3f} MB GPU memory")
return "\n".join(lines)
-def mem_get_info(device: Union[Device, int] = None) -> int:
+def mem_get_info(device: Union[Device, int] = None) -> Tuple[int, int]:
r"""Returns the global free and total GPU memory occupied for a given
device using cudaMemGetInfo.
