diff --git a/torchvision/transforms/v2/functional/_geometry.py b/torchvision/transforms/v2/functional/_geometry.py index 04bb1d2aea7..0c7eab0c04e 100644 --- a/torchvision/transforms/v2/functional/_geometry.py +++ b/torchvision/transforms/v2/functional/_geometry.py @@ -261,7 +261,7 @@ def _do_native_uint8_resize_on_cpu(interpolation: InterpolationMode) -> bool: if torch.compiler.is_compiling(): return True else: - return "AVX2" in torch.backends.cpu.get_cpu_capability() + return torch.backends.cpu.get_cpu_capability() in ("AVX2", "AVX512") return interpolation == InterpolationMode.BICUBIC