expand decorator test to all functional transforms

Author: pmeier

torchvision/prototype/transforms/_presets.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple, cast
 
 import torch
 from torch import Tensor, nn
@@ -41,7 +41,7 @@ def forward(self, img: Tensor) -> Tensor:
             img = F.pil_to_tensor(img)
         img = F.convert_image_dtype(img, torch.float)
         img = F.normalize(img, mean=self._mean, std=self._std)
-        return img
+        return cast(Tensor, img)
 
 
 class Kinect400Eval(nn.Module):
@@ -65,7 +65,7 @@ def forward(self, vid: Tensor) -> Tensor:
         vid = F.resize(vid, self._size, interpolation=self._interpolation)
         vid = F.center_crop(vid, self._crop_size)
         vid = F.convert_image_dtype(vid, torch.float)
-        vid = F.normalize(vid, mean=self._mean, std=self._std)
+        vid = cast(Tensor, F.normalize(vid, mean=self._mean, std=self._std))
         return vid.permute(1, 0, 2, 3)  # (T, C, H, W) => (C, T, H, W)
 
 

torchvision/prototype/transforms/kernels/_geometry.py

@@ -39,7 +39,7 @@ def resize_image(
     new_height, new_width = size
     num_channels, old_height, old_width = image.shape[-3:]
     batch_shape = image.shape[:-3]
-    return _F.resize(
+    return _F.resize(  # type: ignore[no-any-return]
         image.reshape((-1, num_channels, old_height, old_width)),
         size=size,
         interpolation=interpolation,