Fix, not sure why

pytorch · Dec 10, 2024 · b6b863e · b6b863e
1 parent cc76037
commit b6b863e
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/torchvision/io/image.py b/torchvision/io/image.py
@@ -398,9 +398,13 @@ def _load_extra_decoders_once():
 
 def decode_avif(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor:
     _load_extra_decoders_once()
+    if input.dtype != torch.uint8:
+        raise RuntimeError(f"Input tensor must have uint8 data type, got {input.dtype}")
     return torch.ops.extra_decoders_ns.decode_avif(input, mode.value)
 
 
 def decode_heic(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor:
     _load_extra_decoders_once()
+    if input.dtype != torch.uint8:
+        raise RuntimeError(f"Input tensor must have uint8 data type, got {input.dtype}")
     return torch.ops.extra_decoders_ns.decode_heic(input, mode.value)