diff --git a/src/diffusers/utils/testing_utils.py b/src/diffusers/utils/testing_utils.py index 137420945340..e62f245f9ed1 100644 --- 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,