Update GANDLF/models/brain_age.py

Co-authored-by: Sarthak Pati <[email protected]>
mlcommons · Aug 7, 2023 · f189e7b · f189e7b
1 parent 318fc7f
commit f189e7b
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/GANDLF/models/brain_age.py b/GANDLF/models/brain_age.py
@@ -25,9 +25,8 @@ def brainage(parameters):
     try:
         # Load the pretrained VGG16 model
         model = torchvision.models.vgg16(pretrained=True)
-    except Exception as e:
-        print("Error: Failed to load VGG16 model:", str(e))
-        sys.exit(1)
+    except Exception:
+        sys.exit("Error: Failed to load VGG16 model: " + traceback.format_exc())
 
     # Remove the final convolutional layer
     model.final_convolution_layer = None