Unverified Commit e9ae3f80 authored by Kunshang Ji's avatar Kunshang Ji Committed by GitHub
Browse files

[Hardware][XPU] Align memory usage with cuda on xpu (#37029)


Signed-off-by: default avatarKunshang Ji <jikunshang95@gmail.com>
Signed-off-by: default avatarKunshang Ji <kunshang.ji@intel.com>
parent 04cec4f9
...@@ -236,6 +236,7 @@ class XPUPlatform(Platform): ...@@ -236,6 +236,7 @@ class XPUPlatform(Platform):
def get_current_memory_usage( def get_current_memory_usage(
cls, device: torch.types.Device | None = None cls, device: torch.types.Device | None = None
) -> float: ) -> float:
torch.xpu.empty_cache()
torch.xpu.reset_peak_memory_stats(device) torch.xpu.reset_peak_memory_stats(device)
return torch.xpu.max_memory_allocated(device) return torch.xpu.max_memory_allocated(device)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment