O(1) check

Author: NicolasHug

test/test_transforms_v2.py

@@ -4612,7 +4612,7 @@ def test_correctness_image(self, bits, fn):
     def test_error_functional(self, bits):
         with pytest.raises(
             TypeError,
-            match=re.escape(f"bits must be a positive integer in the range [0, 8], get {bits} instead."),
+            match=re.escape(f"bits must be a positive integer in the range [0, 8], got {bits} instead."),
         ):
             F.posterize(make_image(dtype=torch.uint8), bits=bits)
 

torchvision/transforms/v2/functional/_color.py

@@ -460,8 +460,8 @@ def posterize(inpt: torch.Tensor, bits: int) -> torch.Tensor:
 @_register_kernel_internal(posterize, torch.Tensor)
 @_register_kernel_internal(posterize, tv_tensors.Image)
 def posterize_image(image: torch.Tensor, bits: int) -> torch.Tensor:
-    if (not isinstance(bits, int)) or (bits not in range(0, 9)):
-        raise TypeError(f"bits must be a positive integer in the range [0, 8], get {bits} instead.")
+    if not isinstance(bits, int) or not 0 <= bits <= 8:
+        raise TypeError(f"bits must be a positive integer in the range [0, 8], got {bits} instead.")
 
     if image.is_floating_point():
         levels = 1 << bits