Update app.py

app.py

@@ -153,6 +153,9 @@ torch.manual_seed(SEED)
 
 # Parameters
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+USE_CUDA = False
+if device == torch.device("cuda"):
+    USE_CUDA = True
 
 def load_model(root_dir, model_name, model_file_name):
     if CUSTOM_MODEL_FLAG:
@@ -259,7 +262,7 @@ if uploaded_mri_file is not None:
 
     # Build the CAM (Class Activation Map)
     target_layers = [mri_model.model.norm]
-    cam = GradCAM(model=mri_model, target_layers=target_layers, reshape_transform=reshape_transform, use_cuda=True)
+    cam = GradCAM(model=mri_model, target_layers=target_layers, reshape_transform=reshape_transform, use_cuda=USE_CUDA)
     grayscale_cam = cam(input_tensor=image_tensor, targets=[ClassifierOutputTarget(CAM_CLASS_ID)])
     grayscale_cam = grayscale_cam[0, :]