Skip tests on GPU when not available

f-dangel · Oct 19, 2023 · cc4022d · cc4022d
1 parent 96a8e82
commit cc4022d
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/test/optim/test_kfac.py b/test/optim/test_kfac.py
@@ -45,6 +45,8 @@ def test_kfac_single_linear_module(
     Raises:
         AssertionError: If the KFAC approximation is not exact.
     """
+    if not torch.cuda.is_available() and device.type == "cuda":
+        return
     # Fix random seed.
     torch.manual_seed(711)
     # Set up inputs x.
@@ -101,6 +103,8 @@ def test_kfac_deep_linear(
         AssertionError: If the KFAC approximation is not exact for the block
             diagonal.
     """
+    if not torch.cuda.is_available() and device.type == "cuda":
+        return
     # Fix random seed.
     torch.manual_seed(711)
     # Set up inputs x.
@@ -166,6 +170,8 @@ def test_kfac_conv2d_module(
         AssertionError: If the KFAC-reduce approximation is not exact for the
             diagonal or the Conv2d layer or if it is exact for KFAC-expand.
     """
+    if not torch.cuda.is_available() and device.type == "cuda":
+        return
     # Fix random seed.
     torch.manual_seed(711)
     # Set up inputs x.