Update app.py

app.py

@@ -143,17 +143,17 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 ct_root_dir = tempfile.mkdtemp() if CT_MODEL_DIRECTORY is None else CT_MODEL_DIRECTORY
 mri_root_dir = tempfile.mkdtemp() if MRI_MODEL_DIRECTORY is None else MRI_MODEL_DIRECTORY
 
-def load_model(root_dir, model_name, model_file_name, map_location):
+def load_model(root_dir, model_name, model_file_name):
     if CUSTOM_MODEL_FLAG:
         model = Build_Custom_Model(model_name, NUM_CLASSES, pretrained=False).to(device)
     else:
         model = SEResNet50(spatial_dims=2, in_channels=1, num_classes=NUM_CLASSES).to(device)
-    model.load_state_dict(torch.load(os.path.join(root_dir, model_file_name),  map_location=map_location)))
+    model.load_state_dict(torch.load(os.path.join(root_dir, model_file_name),  map_location=device)))
     model.eval()
     return model
 
-ct_model = load_model(ct_root_dir, CT_MODEL_NAME, CT_MODEL_FILE_NAME, device)
-mri_model = load_model(mri_root_dir, MRI_MODEL_NAME, MRI_MODEL_FILE_NAME, device)
+ct_model = load_model(ct_root_dir, CT_MODEL_NAME, CT_MODEL_FILE_NAME)
+mri_model = load_model(mri_root_dir, MRI_MODEL_NAME, MRI_MODEL_FILE_NAME)
 if LIST_MODEL_MODULES:
     for ct_name, _ in ct_model.named_modules():
         print(ct_name)