Compile test fix (#6104)

* update * update

Compile test fix (#6104)
* update * update
d704a730 · Dhruv Nair · GitHub · 49db233b · d704a730
Unverified Commit d704a730 authored Dec 15, 2023 by Dhruv Nair Committed by GitHub Dec 15, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 2 deletions

src/diffusers/utils/testing_utils.py src/diffusers/utils/testing_utils.py +6 -2

No files found.
--- a/src/diffusers/utils/testing_utils.py
+++ b/src/diffusers/utils/testing_utils.py
@@ -820,7 +820,9 @@ def _is_torch_fp16_available(device):

    try:
        x = torch.zeros((2, 2), dtype=torch.float16).to(device)
-        _ = x @ x
+        _ = torch.mul(x, x)
+        return True
+
    except Exception as e:
        if device.type == "cuda":
            raise ValueError(
@@ -838,7 +840,9 @@ def _is_torch_fp64_available(device):

    try:
        x = torch.zeros((2, 2), dtype=torch.float64).to(device)
-        _ = x @ x
+        _ = torch.mul(x, x)
+        return True
+
    except Exception as e:
        if device.type == "cuda":
            raise ValueError(