add scriptable transforms: rgb_to_grayscale

Author: PyExtreme

torchvision/transforms/functional_tensor.py

@@ -49,19 +49,25 @@ def crop(img, top, left, height, width):
 
     return img[..., top:top + height, left:left + width]
 
-def to_grayscale(img, num_output_channels = 3):
+
+def rgb_to_grayscale(img, num_output_channels=3):
     """Convert the given RGB Image Tensor to Grayscale.
 
     Args
         img (Tensor): Image to be converted to Grayscale in the form [C, H, W].
         num_output_channels (int): denotes the number of channels to return after conversion
     Returns:
     Tensor: Grayscale image.
+
+    For RGB to Grayscale conversion, ITU-R 601-2 luma transform is performed which is
+    L = R * 0.2989 + G * 0.5870 + B * 0.1140
     """
     if not F._is_tensor_image(img):
         raise TypeError('tensor is not a torch image.')
 
     if img.size()[0] != 3:
         raise TypeError('Input Image does not contain 3 Channels')
 
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    img = img.to(device)
     return (0.2989 * img[0] + 0.5870 * img[1] + 0.1140 * img[2]).to(img.dtype)