Update

app.py

@@ -121,7 +121,7 @@ def visualize_attention(
     attentions_for_rollout = []
     for layer_name, attn_map in attention_maps.items():
         print(f"Attention map shape for {layer_name}: {attn_map.shape}")
-        attn_map = attn_map[0]  # Remove batch dimension
+        attn_map = attn_map[0].detach()  # Remove batch dimension and detach
 
         attentions_for_rollout.append(attn_map)
 
@@ -148,7 +148,7 @@ def visualize_attention(
         # Interpolate to match image size
         attn_map = attn_map.unsqueeze(0).unsqueeze(0)
         attn_map = F.interpolate(attn_map, size=(image_np.shape[0], image_np.shape[1]), mode='bilinear', align_corners=False)
-        attn_map = attn_map.squeeze().cpu().numpy()  # Move to CPU before converting to numpy
+        attn_map = attn_map.squeeze().cpu().detach().numpy()  # Move to CPU, detach, and convert to numpy
 
         # Normalize attention map
         attn_map = (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min())
@@ -177,8 +177,8 @@ def visualize_attention(
         visualizations.append(vis_image)
         plt.close(fig)
 
-    # Ensure tensors are on CPU before converting to numpy
-    attentions_for_rollout = [attn.cpu() for attn in attentions_for_rollout]
+    # Ensure tensors are on CPU and detached before converting to numpy
+    attentions_for_rollout = [attn.cpu().detach() for attn in attentions_for_rollout]
 
     # Calculate rollout
     rollout_mask = rollout(attentions_for_rollout, discard_ratio, head_fusion, num_prefix_tokens)