disable slow test in ci

quintenroets · Apr 9, 2024 · 0dbfa8e · 0dbfa8e
1 parent 548cc16
commit 0dbfa8e
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -1,5 +1,12 @@
+import pytest
+import torch
 from revnets.main.main import main
 
+gpu_available = torch.cuda.is_available()
 
+
+@pytest.mark.skipif(
+    not gpu_available, reason="Only test model training if GPU is available"
+)
 def test_main() -> None:
     main()