diff --git a/models/efficientnet.py b/models/efficientnet.py
index fb196fc..e644a3d 100644
--- a/models/efficientnet.py
+++ b/models/efficientnet.py
@@ -10,16 +10,20 @@ def efficientnet(input_shape, n_class=2, variant='B0', **kwargs):
         'pooling': None,
         'classes': n_class,
     }
-    if len(input_shape) == 4:
-        effnet_layer = getattr(efn, f'EfficientNet{variant}')(**default_effnet_params)
-        modal = '3D'
-    elif len(input_shape) == 3:
-        effnet_layer = getattr(keras.applications, f'EfficientNet{variant}')(
-            classifier_activation='softmax', **default_effnet_params
-        )
-        modal = '2D'
-    else:
-        raise ValueError('input_shape is expected as an array ranked 3 or 4')
+    try:
+        if len(input_shape) == 4:
+            effnet_layer = getattr(efn, f'EfficientNet{variant}')(**default_effnet_params)
+            modal = '3D'
+        elif len(input_shape) == 3:
+            effnet_layer = getattr(keras.applications, f'EfficientNet{variant}')(
+                classifier_activation='softmax', **default_effnet_params
+            )
+            modal = '2D'
+        else:
+            raise ValueError('input_shape is expected as an array ranked 3 or 4')
+    except AttributeError:
+        print(f'No EfficientNet variation found for {variant}')
+        raise SystemExit
 
     inputs = keras.Input(input_shape)
     outputs = effnet_layer(inputs)