map BACKEND_RESET_MAX_MEMORY_ALLOCATED to reset_peak_memory_stats on XPU

yao-matrix · yao-matrix · commit b021a7697bfb · 2025-04-01T20:14:28.000-07:00
Signed-off-by: YAO Matrix &lt;matrix.yao@intel.com&gt;
diff --git a/src/diffusers/utils/testing_utils.py b/src/diffusers/utils/testing_utils.py
@@ -1161,7 +1161,7 @@ def _is_torch_fp64_available(device):
     }
     BACKEND_RESET_MAX_MEMORY_ALLOCATED = {
         "cuda": torch.cuda.reset_max_memory_allocated,
-        "xpu": None,
+        "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None),
         "cpu": None,
         "mps": None,
         "default": None,

Original file line number	Diff line number	Diff line change
`@@ -1161,7 +1161,7 @@ def _is_torch_fp64_available(device):`
`1161`	`1161`	`}`
`1162`	`1162`	`BACKEND_RESET_MAX_MEMORY_ALLOCATED = {`
`1163`	`1163`	`"cuda": torch.cuda.reset_max_memory_allocated,`
`1164`		`- "xpu": None,`
	`1164`	`+ "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None),`
`1165`	`1165`	`"cpu": None,`
`1166`	`1166`	`"mps": None,`
`1167`	`1167`	`"default": None,`