Add error handling for brain age model

mlcommons · Aug 6, 2023 · 9d56e70 · 9d56e70
1 parent f189e7b
commit 9d56e70
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/GANDLF/models/brain_age.py b/GANDLF/models/brain_age.py
@@ -18,9 +18,7 @@ def brainage(parameters):
     """
 
     # Check that the input data is 2D
-    if parameters["model"]["dimension"] != 2:
-        print("Error: Brain Age predictions only work on 2D data")
-        sys.exit(1)
+    assert parameters["model"]["dimension"] == 2, "Brain Age predictions only work on 2D data"
 
     try:
         # Load the pretrained VGG16 model