Use device instead of cuda

app.py

@@ -55,7 +55,7 @@ def get_attention_mask(img_tensor):
 
     # Average the attention at each layer over all heads
     attention_probs = torch.mean(attention_probs, dim=1)
-    residual = torch.eye(attention_probs.size(-1)).to("cuda")
+    residual = torch.eye(attention_probs.size(-1)).to(device)
     attention_probs = 0.5 * attention_probs + 0.5 * residual
 
     # normalize by layer