pytest mps check fix

jeshraghian · Oct 10, 2023 · fc383f9 · fc383f9
1 parent 71b14d6
commit fc383f9
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/tests/test_backprop.py b/tests/test_backprop.py
@@ -9,8 +9,14 @@
 from unittest import mock
 import torch
 
-device = (
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
+device = device = (
+    torch.device("cuda")
+    if torch.cuda.is_available()
+    else (
+        torch.device("mps")
+        if torch.backends.mps.is_available()
+        else torch.device("cpu")
+    )
 )