skip testing big models

neuronets · Mar 19, 2024 · 4d67696 · 4d67696
1 parent 91f74e4
commit 4d67696
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/nobrainer/models/attention_unet.py b/nobrainer/models/attention_unet.py
@@ -1,6 +1,6 @@
 """Model definition for Attention U-Net.
 Adapted from https://github.com/nikhilroxtomar/Semantic-Segmentation-Architecture/blob/main/TensorFlow/attention-unet.py
-"""
+"""  # noqa: E501
 
 from tensorflow.keras import layers
 import tensorflow.keras.layers as L

diff --git a/nobrainer/models/tests/models_test.py b/nobrainer/models/tests/models_test.py
@@ -1,3 +1,5 @@
+import os
+
 import numpy as np
 import pytest
 import tensorflow as tf
@@ -21,6 +23,8 @@
 from ..vnet import vnet
 from ..vox2vox import Vox_ensembler, vox_gan
 
+IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true"
+
 
 def model_test(model_cls, n_classes, input_shape, kwds={}):
     """Tests for models."""
@@ -257,9 +261,11 @@ def test_attention_unet_with_inception():
     )
 
 
+@pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Cannot test in GitHub Actions")
 def test_unetr():
     model_test(unetr, n_classes=1, input_shape=(1, 96, 96, 96, 1))
 
 
+@pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Cannot test in GitHub Actions")
 def test_variational_meshnet():
     model_test(variational_meshnet, n_classes=1, input_shape=(1, 128, 128, 128, 1))