fix elastic float16 output

Author: pmeier

test/test_transforms_v2_refactored.py

@@ -2491,7 +2491,7 @@ def _make_displacement(self, inpt):
         interpolation=[transforms.InterpolationMode.NEAREST, transforms.InterpolationMode.BILINEAR],
         fill=EXHAUSTIVE_TYPE_FILLS,
     )
-    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8, torch.float16])
     @pytest.mark.parametrize("device", cpu_and_cuda())
     def test_kernel_image(self, param, value, dtype, device):
         image = make_image_tensor(dtype=dtype, device=device)
@@ -2502,6 +2502,7 @@ def test_kernel_image(self, param, value, dtype, device):
             displacement=self._make_displacement(image),
             **{param: value},
             check_scripted_vs_eager=not (param == "fill" and isinstance(value, (int, float))),
+            check_cuda_vs_cpu=dtype is not torch.float16,
         )
 
     @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))

torchvision/transforms/v2/functional/_geometry.py

@@ -1743,7 +1743,12 @@ def elastic_image(
     grid = _create_identity_grid((height, width), device=device, dtype=dtype).add_(
         displacement.to(dtype=dtype, device=device)
     )
-    return _apply_grid_transform(image, grid, interpolation.value, fill=fill)
+    output = _apply_grid_transform(image, grid, interpolation.value, fill=fill)
+
+    if is_cpu_half:
+        output = output.to(torch.float16)
+
+    return output
 
 
 @_register_kernel_internal(elastic, PIL.Image.Image)