resolved #289

neuronets · Mar 12, 2024 · bf224f5 · bf224f5
1 parent eb6d11e
commit bf224f5
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/nobrainer/models/tests/models_test.py b/nobrainer/models/tests/models_test.py
@@ -16,6 +16,7 @@
 from ..progressivegan import progressivegan
 from ..unet import unet
 from ..unet_lstm import unet_lstm
+from ..unetr import unetr
 from ..vnet import vnet
 from ..vox2vox import Vox_ensembler, vox_gan
 
@@ -253,3 +254,7 @@ def test_attention_unet_with_inception():
     model_test(
         attention_unet_with_inception, n_classes=1, input_shape=(1, 64, 64, 64, 1)
     )
+
+
+def test_unetr():
+    model_test(unetr, n_classes=1, input_shape=(1, 96, 96, 96, 1))