get_memory_info can handle default arg (#7590)

pytorch · Jul 1, 2024 · 020f3f5 · 020f3f5
1 parent 6873d33
commit 020f3f5
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/test/pjrt/test_runtime_tpu.py b/test/pjrt/test_runtime_tpu.py
@@ -255,7 +255,7 @@ def test_execute_time_metric(self):
 
   @staticmethod
   def _memory_usage():
-    return xm.get_memory_info(torch_xla.device())
+    return xm.get_memory_info()
 
   def test_memory_usage(self):
     results = pjrt.run_multiprocess(self._memory_usage)

diff --git a/torch_xla/core/xla_model.py b/torch_xla/core/xla_model.py
@@ -1488,15 +1488,18 @@ class MemoryInfo(TypedDict):
   bytes_limit: int
 
 
-def get_memory_info(device: torch.device) -> MemoryInfo:
+def get_memory_info(device: Optional[torch.device] = None) -> MemoryInfo:
   """Retrieves the device memory usage.
 
   Args:
-    device: The device whose memory information are requested.
+    device: Optional[torch.device] The device whose memory information are requested.
+    If not passed will use the default device.
 
   Returns:
     MemoryInfo dict with memory usage for the given device.
   """
+  if device == None:
+    device = xla_device()
   return torch_xla._XLAC._xla_memory_info(str(device))