From 0dbfa8e1090beb77d161718e4fbe63d365d38338 Mon Sep 17 00:00:00 2001
From: Quinten <quinten.roets@gmail.com>
Date: Mon, 8 Apr 2024 22:54:45 -0700
Subject: [PATCH] disable slow test in ci

---
 tests/test_main.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tests/test_main.py b/tests/test_main.py
index 8b6afb4..2d32753 100644
--- 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()