Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
This commit is contained in:
Wang, Yi A 2025-03-12 00:47:02 -07:00
parent 468da545af
commit 9997047d8a

View File

@ -20,7 +20,9 @@ def get_cuda_free_memory(device, memory_fraction):
def get_xpu_free_memory(device, memory_fraction):
total_free_memory, total_xpu_memory = torch.xpu.mem_get_info(device)
memory_fraction = float(os.getenv("XPU_MEMORY_FRACTION", "1.0"))
free_memory = max(0, total_free_memory - (1 - memory_fraction) * total_xpu_memory)
free_memory = max(
0, int(total_free_memory - (1 - memory_fraction) * total_xpu_memory)
)
return free_memory