Skip to content

Commit 3db0b12

Browse files
authored
Allow native uint8 resize on avx512 (#9190)
1 parent 4bc406d commit 3db0b12

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

torchvision/transforms/v2/functional/_geometry.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -261,7 +261,7 @@ def _do_native_uint8_resize_on_cpu(interpolation: InterpolationMode) -> bool:
261261
if torch.compiler.is_compiling():
262262
return True
263263
else:
264-
return "AVX2" in torch.backends.cpu.get_cpu_capability()
264+
return torch.backends.cpu.get_cpu_capability() in ("AVX2", "AVX512")
265265

266266
return interpolation == InterpolationMode.BICUBIC
267267

0 commit comments

Comments
 (0)