cuda skip

Author: SkafteNicki

tests/test_nnj.py

@@ -149,6 +149,9 @@ def test_jacobians(self, model, input_shape, device, dtype):
     @pytest.mark.parametrize("return_jac", [True, False])
     def test_jac_return(self, model, input_shape, device, return_jac):
         """ Test that all models returns the jacobian output if asked for it """
+        if device == "cuda" and not torch.cuda.is_available():
+            pytest.skip("Test requires cuda support")
+
         input = torch.randn(*input_shape, device=device)
         model = deepcopy(model).to(device)
         output = model(input, jacobian=return_jac)