Minor perf on saturation uint8 (#6940)

0cc90808 · Vasilis Vryniotis · GitHub · 98f1a409 · 0cc90808
Unverified Commit 0cc90808 authored Nov 10, 2022 by Vasilis Vryniotis Committed by GitHub Nov 10, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

torchvision/prototype/transforms/functional/_color.py torchvision/prototype/transforms/functional/_color.py +5 -1

No files found.
--- a/torchvision/prototype/transforms/functional/_color.py
+++ b/torchvision/prototype/transforms/functional/_color.py
@@ -53,7 +53,11 @@ def adjust_saturation_image_tensor(image: torch.Tensor, saturation_factor: float
    if c == 1:  # Match PIL behaviour
        return image

-    return _blend(image, _rgb_to_gray(image), saturation_factor)
+    grayscale_image = _rgb_to_gray(image, cast=False)
+    if not image.is_floating_point():
+        grayscale_image = grayscale_image.floor_()
+
+    return _blend(image, grayscale_image, saturation_factor)


 adjust_saturation_image_pil = _FP.adjust_saturation