From 07c1323ef4f4d1033140e6a153a1a040b763afa4 Mon Sep 17 00:00:00 2001 From: Vasilis Vryniotis Date: Thu, 10 Nov 2022 11:28:34 +0000 Subject: [PATCH] Minor perf on saturation uint8 --- torchvision/prototype/transforms/functional/_color.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/torchvision/prototype/transforms/functional/_color.py b/torchvision/prototype/transforms/functional/_color.py index 671f3fcf90f..376021f0974 100644 --- 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