Update app.py

app.py

@@ -14,8 +14,8 @@ from pytorch_grad_cam.utils.image import show_cam_on_image
 with zipfile.ZipFile("examples.zip","r") as zip_ref:
     zip_ref.extractall(".")
 
-DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
-
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+'''cuda:0'''
 mtcnn = MTCNN(
     select_largest=False,
     post_process=False,
@@ -66,8 +66,8 @@ def predict(input_image:Image.Image, true_label:str):
 
     target_layers=[model.block8.branch1[-1]]
     use_cuda = True if torch.cuda.is_available() else False
-    print ("Cuda :: "+ use_cuda)
-    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=use_cuda)
+    '''print ("Cuda :: ", use_cuda)'''
+    cam = GradCAM(model=model, target_layers=target_layers) ''', use_cuda=use_cuda)'''
     targets = [ClassifierOutputTarget(0)]
 
     grayscale_cam = cam(input_tensor=face, targets=targets, eigen_smooth=True)