remove cuda

generic_utils.py

@@ -24,10 +24,10 @@ def show_cam_on_image(img, mask):
 
 
 # initialize ViT pretrained
-model = vit_LRP(pretrained=True).cuda()
+model = vit_LRP(pretrained=True)
 model.eval()
 attribution_generator = LRP(model)
-model_baseline = vit(pretrained=True).cuda()
+model_baseline = vit(pretrained=True)
 model_baseline.eval()
 baselines_generator = Baselines(model_baseline)
 
@@ -37,16 +37,16 @@ def generate_visualization(
 ):
     if LRP:
         transformer_attribution = attribution_generator.generate_LRP(
-            original_image.unsqueeze(0).cuda(), method=method, index=class_index
+            original_image.unsqueeze(0), method=method, index=class_index
         ).detach()
     else:
         if method == "gradcam":
             transformer_attribution = baselines_generator.generate_cam_attn(
-                original_image.unsqueeze(0).cuda(), index=class_index
+                original_image.unsqueeze(0), index=class_index
             ).detach()
         else:
             transformer_attribution = baselines_generator.generate_rollout(
-                original_image.unsqueeze(0).cuda()
+                original_image.unsqueeze(0)
             ).detach()
     if method != "full":
         transformer_attribution = transformer_attribution.reshape(1, 1, 14, 14)