Compile test fix (huggingface#6104)

* update * update
AmericanPresidentJimmyCarter · Apr 26, 2024 · a8221c5 · a8221c5
1 parent 86087cd
commit a8221c5
Showing 1 changed file with 6 additions and 2 deletions.
diff --git 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(