Update test_utils.py

makgyver · Oct 5, 2024 · 04a5922 · 04a5922
1 parent e7559f4
commit 04a5922
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -488,7 +488,9 @@ def test_plot_dist(mock_show):
 
 def test_check_mem():
     net = MNIST_2NN()
-    assert check_model_fit_mem(net, (28 * 28,), 100, "mps", True)
+
+    if torch.backends.mps.is_available():
+        assert check_model_fit_mem(net, (28 * 28,), 100, "mps", True)
 
     if torch.cuda.is_available():
         assert check_model_fit_mem(net, (28 * 28,), 100, "cuda")