fixing image_encoder to work with cuda_graphs

Summary: the combination of tensors on multiple devices in get_rel_pos
was preventing cuda graphs from correctly optimizing things

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 0fea0e19e5bf0ee44a19669fe33e7e16002a55af
Pull Request resolved: #393

segment_anything/modeling/image_encoder.py

@@ -315,8 +315,8 @@ def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor
         rel_pos_resized = rel_pos
 
     # Scale the coords with short length if shapes for q and k are different.
-    q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
-    k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
+    q_coords = (torch.arange(q_size).to(rel_pos.device)[:, None] * max(k_size / q_size, 1.0))
+    k_coords = (torch.arange(k_size).to(rel_pos.device)[None, :] * max(q_size / k_size, 1.0))
     relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)
 
     return rel_pos_resized[relative_coords.long()]