Moving tensors to the right device (#3870)

* Moving tensors to the right device.

* Switch to gpu.medium

.circleci/config.yml

@@ -661,7 +661,7 @@ jobs:
     <<: *binary_common
     machine:
       image: ubuntu-1604-cuda-10.2:202012-01
-    resource_class: gpu.small
+    resource_class: gpu.nvidia.medium
     environment:
       image_name: "pytorch/manylinux-cuda102"
       PYTHON_VERSION: << parameters.python_version >>

.circleci/config.yml.in

@@ -661,7 +661,7 @@ jobs:
     <<: *binary_common
     machine:
       image: ubuntu-1604-cuda-10.2:202012-01
-    resource_class: gpu.small
+    resource_class: gpu.nvidia.medium
     environment:
       image_name: "pytorch/manylinux-cuda102"
       PYTHON_VERSION: << parameters.python_version >>

torchvision/models/detection/roi_heads.py

@@ -274,7 +274,7 @@ def heatmaps_to_keypoints(maps, rois):
         xy_preds[i, 0, :] = x + offset_x[i]
         xy_preds[i, 1, :] = y + offset_y[i]
         xy_preds[i, 2, :] = 1
-        end_scores[i, :] = roi_map[torch.arange(num_keypoints), y_int, x_int]
+        end_scores[i, :] = roi_map[torch.arange(num_keypoints, device=roi_map.device), y_int, x_int]
 
     return xy_preds.permute(0, 2, 1), end_scores
 

torchvision/transforms/functional_tensor.py

@@ -391,12 +391,12 @@ def _pad_symmetric(img: Tensor, padding: List[int]) -> Tensor:
     x_indices = [i for i in range(in_sizes[-1])]  # [0, 1, 2, 3, ...]
     left_indices = [i for i in range(padding[0] - 1, -1, -1)]  # e.g. [3, 2, 1, 0]
     right_indices = [-(i + 1) for i in range(padding[1])]  # e.g. [-1, -2, -3]
-    x_indices = torch.tensor(left_indices + x_indices + right_indices)
+    x_indices = torch.tensor(left_indices + x_indices + right_indices, device=img.device)
 
     y_indices = [i for i in range(in_sizes[-2])]
     top_indices = [i for i in range(padding[2] - 1, -1, -1)]
     bottom_indices = [-(i + 1) for i in range(padding[3])]
-    y_indices = torch.tensor(top_indices + y_indices + bottom_indices)
+    y_indices = torch.tensor(top_indices + y_indices + bottom_indices, device=img.device)
 
     ndim = img.ndim
     if ndim == 3: