use torch.where over boolean masking

[pmeier]

This solves the graph break for torch.compile that stems from pytorch/pytorch#114220. The fix proposed there is no fix for now, but rather just a pointer to setting a feature flag. This flag is supposed to become the default in the future. Still, using torch.where here adds no complexity and will be faster on CUDA as well.

cc @vfdev-5

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/8171

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (32 Unrelated Failures)

As of commit a3d9eb3 with merge base 6640e49 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• CMake / linux (linux.12xlarge, cpu) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• CMake / linux (linux.g5.4xlarge.nvidia.gpu, cuda, 11.8) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• CMake / macos (macos-12) / macos-job (gh)
  The process '/usr/local/bin/git' failed with exit code 128
• CMake / macos (macos-m1-12) / macos-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• CMake / windows (windows.4xlarge, cpu) / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128
• CMake / windows (windows.g5.4xlarge.nvidia.gpu, cuda, 11.8) / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128
• Docs / build / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Lint / c-source / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Lint / python-source-and-configs / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Lint / python-types / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Prototype tests on Linux / unittests-prototype (3.10, linux.12xlarge, cpu) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Prototype tests on Linux / unittests-prototype (3.11, linux.12xlarge, cpu) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Prototype tests on Linux / unittests-prototype (3.8, linux.12xlarge, cpu) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Prototype tests on Linux / unittests-prototype (3.8, linux.g5.4xlarge.nvidia.gpu, cuda, 11.8) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Prototype tests on Linux / unittests-prototype (3.9, linux.12xlarge, cpu) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Tests / onnx / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-extended / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-linux (3.10, linux.12xlarge, cpu) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-linux (3.11, linux.12xlarge, cpu) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-linux (3.8, linux.12xlarge, cpu) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-linux (3.8, linux.g5.4xlarge.nvidia.gpu, cuda, 11.8) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-linux (3.9, linux.12xlarge, cpu) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-macos (3.10, macos-12) / macos-job (gh)
  The process '/usr/local/bin/git' failed with exit code 128
• Tests / unittests-macos (3.11, macos-12) / macos-job (gh)
  The process '/usr/local/bin/git' failed with exit code 128
• Tests / unittests-macos (3.8, macos-12) / macos-job (gh)
  The process '/usr/local/bin/git' failed with exit code 128
• Tests / unittests-macos (3.8, macos-m1-12) / macos-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-macos (3.9, macos-12) / macos-job (gh)
  The process '/usr/local/bin/git' failed with exit code 128
• Tests / unittests-windows (3.10, windows.4xlarge, cpu) / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128
• Tests / unittests-windows (3.11, windows.4xlarge, cpu) / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128
• Tests / unittests-windows (3.8, windows.4xlarge, cpu) / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128
• Tests / unittests-windows (3.8, windows.g5.4xlarge.nvidia.gpu, cuda, 11.8) / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128
• Tests / unittests-windows (3.9, windows.4xlarge, cpu) / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[lezcano]

Also note that this will also generate better code even when pytorch/pytorch#114220 is set as the default.

[vfdev-5]

LGTM, thanks @pmeier !

[github-actions]

Hey @pmeier!

You merged this PR, but no labels were added. The list of valid labels is available at https://github.com/pytorch/vision/blob/main/.github/process_commit.py