Update

configs.py

@@ -44,8 +44,8 @@ LEARNING_RATE = 0.0001
 STEP_SIZE = 10
 GAMMA = 0.3
 CUTMIX_ALPHA = 0.3
-# DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-DEVICE = torch.device("cpu")
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+# DEVICE = torch.device("cpu")
 NUM_PRINT = 100
 TASK = 1
 WARMUP_EPOCHS = 5

extract.py

@@ -59,7 +59,7 @@ def extract_gradcam(image_path=None, save_path=None):
             input_tensor = input_tensor.to(DEVICE)
 
             # Create a GradCAMPlusPlus object
-            cam = GradCAMPlusPlus(model=model, target_layers=[target_layer], use_cuda=True)
+            cam = GradCAMPlusPlus(model=model, target_layers=[target_layer])
 
             # Generate the GradCAM heatmap
             grayscale_cam = cam(input_tensor=input_tensor)[0]