[Platform] move current_memory_usage() into platform (#11369)

Signed-off-by: Shanshan Shen <467638484@qq.com>
This commit is contained in:
Shanshan Shen
2025-01-15 11:38:25 +08:00
committed by GitHub
parent 1a51b9f872
commit 9ddac56311
5 changed files with 31 additions and 7 deletions

View File

@@ -94,3 +94,10 @@ class XPUPlatform(Platform):
def is_pin_memory_available(cls):
logger.warning("Pin memory is not supported on XPU.")
return False
@classmethod
def get_current_memory_usage(cls,
device: Optional[torch.types.Device] = None
) -> float:
torch.xpu.reset_peak_memory_stats(device)
return torch.xpu.max_memory_allocated(device)