Added cuda and torch.jit tests for rgb-hsv tensor conversion ops (#2639)

* Added cuda torch.jit tests for rgb-hsv tensor conversion ops

* Reverted deleted test

test/test_functional_tensor.py

@@ -67,40 +67,46 @@ class Tester(TransformsTester):
             self.compareTensorToPIL(img_tensor_cropped, pil_img_cropped)
 
     def test_hsv2rgb(self):
+        scripted_fn = torch.jit.script(F_t._hsv2rgb)
         shape = (3, 100, 150)
-        for _ in range(20):
-            img = torch.rand(*shape, dtype=torch.float)
-            ft_img = F_t._hsv2rgb(img).permute(1, 2, 0).flatten(0, 1)
+        for _ in range(10):
+            hsv_img = torch.rand(*shape, dtype=torch.float, device=self.device)
+            rgb_img = F_t._hsv2rgb(hsv_img)
+            ft_img = rgb_img.permute(1, 2, 0).flatten(0, 1)
 
-            h, s, v, = img.unbind(0)
-            h = h.flatten().numpy()
-            s = s.flatten().numpy()
-            v = v.flatten().numpy()
+            h, s, v, = hsv_img.unbind(0)
+            h = h.flatten().cpu().numpy()
+            s = s.flatten().cpu().numpy()
+            v = v.flatten().cpu().numpy()
 
             rgb = []
             for h1, s1, v1 in zip(h, s, v):
                 rgb.append(colorsys.hsv_to_rgb(h1, s1, v1))
-
-            colorsys_img = torch.tensor(rgb, dtype=torch.float32)
+            colorsys_img = torch.tensor(rgb, dtype=torch.float32, device=self.device)
             max_diff = (ft_img - colorsys_img).abs().max()
             self.assertLess(max_diff, 1e-5)
 
+            s_rgb_img = scripted_fn(hsv_img)
+            self.assertTrue(rgb_img.allclose(s_rgb_img))
+
     def test_rgb2hsv(self):
+        scripted_fn = torch.jit.script(F_t._rgb2hsv)
         shape = (3, 150, 100)
-        for _ in range(20):
-            img = torch.rand(*shape, dtype=torch.float)
-            ft_hsv_img = F_t._rgb2hsv(img).permute(1, 2, 0).flatten(0, 1)
+        for _ in range(10):
+            rgb_img = torch.rand(*shape, dtype=torch.float, device=self.device)
+            hsv_img = F_t._rgb2hsv(rgb_img)
+            ft_hsv_img = hsv_img.permute(1, 2, 0).flatten(0, 1)
 
-            r, g, b, = img.unbind(0)
-            r = r.flatten().numpy()
-            g = g.flatten().numpy()
-            b = b.flatten().numpy()
+            r, g, b, = rgb_img.unbind(0)
+            r = r.flatten().cpu().numpy()
+            g = g.flatten().cpu().numpy()
+            b = b.flatten().cpu().numpy()
 
             hsv = []
             for r1, g1, b1 in zip(r, g, b):
                 hsv.append(colorsys.rgb_to_hsv(r1, g1, b1))
 
-            colorsys_img = torch.tensor(hsv, dtype=torch.float32)
+            colorsys_img = torch.tensor(hsv, dtype=torch.float32, device=self.device)
 
             ft_hsv_img_h, ft_hsv_img_sv = torch.split(ft_hsv_img, [1, 2], dim=1)
             colorsys_img_h, colorsys_img_sv = torch.split(colorsys_img, [1, 2], dim=1)
@@ -108,9 +114,11 @@ class Tester(TransformsTester):
             max_diff_h = ((colorsys_img_h * 2 * math.pi).sin() - (ft_hsv_img_h * 2 * math.pi).sin()).abs().max()
             max_diff_sv = (colorsys_img_sv - ft_hsv_img_sv).abs().max()
             max_diff = max(max_diff_h, max_diff_sv)
-
             self.assertLess(max_diff, 1e-5)
 
+            s_hsv_img = scripted_fn(rgb_img)
+            self.assertTrue(hsv_img.allclose(s_hsv_img))
+
     def test_rgb_to_grayscale(self):
         script_rgb_to_grayscale = torch.jit.script(F.rgb_to_grayscale)
 

torchvision/transforms/functional_tensor.py

@@ -423,12 +423,13 @@ def _rgb2hsv(img):
 
     cr = maxc - minc
     # Since `eqc => cr = 0`, replacing denominator with 1 when `eqc` is fine.
-    s = cr / torch.where(eqc, maxc.new_ones(()), maxc)
+    ones = torch.ones_like(maxc)
+    s = cr / torch.where(eqc, ones, maxc)
     # Note that `eqc => maxc = minc = r = g = b`. So the following calculation
     # of `h` would reduce to `bc - gc + 2 + rc - bc + 4 + rc - bc = 6` so it
     # would not matter what values `rc`, `gc`, and `bc` have here, and thus
     # replacing denominator with 1 when `eqc` is fine.
-    cr_divisor = torch.where(eqc, maxc.new_ones(()), cr)
+    cr_divisor = torch.where(eqc, ones, cr)
     rc = (maxc - r) / cr_divisor
     gc = (maxc - g) / cr_divisor
     bc = (maxc - b) / cr_divisor
@@ -452,7 +453,7 @@ def _hsv2rgb(img):
     t = torch.clamp((v * (1.0 - s * (1.0 - f))), 0.0, 1.0)
     i = i % 6
 
-    mask = i == torch.arange(6)[:, None, None]
+    mask = i == torch.arange(6, device=i.device)[:, None, None]
 
     a1 = torch.stack((v, q, p, p, t, v))
     a2 = torch.stack((t, v, v, q, p, p))