Merge branch 'main' into arm64_enablement

Author: alinpahontu2912

.github/workflows/build-wheels-aarch64-linux.yml

@@ -25,7 +25,7 @@ jobs:
       os: linux-aarch64
       test-infra-repository: pytorch/test-infra
       test-infra-ref: main
-      with-cuda: disable
+      with-cuda: enable
   build:
     needs: generate-matrix
     strategy:

packaging/pre_build_script_arm64.sh

@@ -13,8 +13,6 @@ export TRIPLET_FILE="triplets/arm64-windows.cmake"
 export PYTORCH_VERSION="$PYTORCH_VERSION"
 export CHANNEL="$CHANNEL"
 
-echo "channel: $CHANNEL"
-
 # Dependencies
 mkdir -p "$DOWNLOADS_DIR"
 mkdir -p "$DEPENDENCIES_DIR"
@@ -27,15 +25,17 @@ git clone https://github.com/microsoft/vcpkg.git
 cd vcpkg || exit
 ./bootstrap-vcpkg.sh
 
-# # Set vcpkg to only build release packages
+# Set vcpkg to only build release packages
+
 echo "set(VCPKG_BUILD_TYPE release)" >> "$TRIPLET_FILE"
 
 # Install dependencies using vcpkg
 ./vcpkg install libjpeg-turbo:arm64-windows --x-install-root="$DEPENDENCIES_DIR"
 ./vcpkg install libwebp:arm64-windows --x-install-root="$DEPENDENCIES_DIR"
 ./vcpkg install libpng[tools]:arm64-windows --x-install-root="$DEPENDENCIES_DIR"
 
-# Copy files using cp (replace robocopy)
+# Copy files using cp
+
 cp "$DEPENDENCIES_DIR/arm64-windows/lib/libpng16.lib" "$DEPENDENCIES_DIR/arm64-windows/lib/libpng.lib"
 cp "$DEPENDENCIES_DIR/arm64-windows/bin/libpng16.dll" "$DEPENDENCIES_DIR/arm64-windows/bin/libpng.dll"
 cp "$DEPENDENCIES_DIR/arm64-windows/bin/libpng16.pdb" "$DEPENDENCIES_DIR/arm64-windows/bin/libpng.pdb"
@@ -62,7 +62,7 @@ if [ "$CHANNEL" = "release" ]; then
   pip3 install --pre torch --index-url https://download.pytorch.org/whl/torch/
 elif [ "$CHANNEL" = "test" ]; then
   echo "Installing PyTorch version $PYTORCH_VERSION."
-  pip3 install torch=="$PYTORCH_VERSION"
+  pip3 install --pre torch=="$PYTORCH_VERSION" --index-url https://download.pytorch.org/whl/test
 else
   echo "CHANNEL is not set, installing PyTorch from nightly."
   pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu

torchvision/models/detection/anchor_utils.py

@@ -74,16 +74,20 @@ def generate_anchors(
         return base_anchors.round()
 
     def set_cell_anchors(self, dtype: torch.dtype, device: torch.device):
-        self.cell_anchors = [cell_anchor.to(dtype=dtype, device=device) for cell_anchor in self.cell_anchors]
+        return [cell_anchor.to(dtype=dtype, device=device) for cell_anchor in self.cell_anchors]
 
     def num_anchors_per_location(self) -> list[int]:
         return [len(s) * len(a) for s, a in zip(self.sizes, self.aspect_ratios)]
 
     # For every combination of (a, (g, s), i) in (self.cell_anchors, zip(grid_sizes, strides), 0:2),
     # output g[i] anchors that are s[i] distance apart in direction i, with the same dimensions as a.
-    def grid_anchors(self, grid_sizes: list[list[int]], strides: list[list[Tensor]]) -> list[Tensor]:
+    def grid_anchors(
+        self,
+        grid_sizes: list[list[int]],
+        strides: list[list[Tensor]],
+        cell_anchors: list[torch.Tensor],
+    ) -> list[Tensor]:
         anchors = []
-        cell_anchors = self.cell_anchors
         torch._assert(cell_anchors is not None, "cell_anchors should not be None")
         torch._assert(
             len(grid_sizes) == len(strides) == len(cell_anchors),
@@ -123,8 +127,8 @@ def forward(self, image_list: ImageList, feature_maps: list[Tensor]) -> list[Ten
             ]
             for g in grid_sizes
         ]
-        self.set_cell_anchors(dtype, device)
-        anchors_over_all_feature_maps = self.grid_anchors(grid_sizes, strides)
+        cell_anchors = self.set_cell_anchors(dtype, device)
+        anchors_over_all_feature_maps = self.grid_anchors(grid_sizes, strides, cell_anchors)
         anchors: list[list[torch.Tensor]] = []
         for _ in range(len(image_list.image_sizes)):
             anchors_in_image = [anchors_per_feature_map for anchors_per_feature_map in anchors_over_all_feature_maps]