Update ToDtype to avoid unnecessary to() calls and fixing types on Transform (#6773)

* Fix `ToDtype` to avoid errors when a type is not defined.

* Nit `(features.is_simple_tensor, features._Feature)` to `(Tensor,)`

* Fixing linter

* Adding comment.

* Switch back to indexing. Python's default dict seems to have a nasty behaviour.

Author: datumbox

torchvision/prototype/transforms/_misc.py

@@ -157,7 +157,10 @@ def __init__(self, dtype: Union[torch.dtype, Dict[Type, torch.dtype]]) -> None:
         self.dtype = dtype
 
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
-        return inpt.to(self.dtype[type(inpt)])
+        dtype = self.dtype[type(inpt)]
+        if dtype is None:
+            return inpt
+        return inpt.to(dtype=dtype)
 
 
 class RemoveSmallBoundingBoxes(Transform):

torchvision/prototype/transforms/_transform.py

@@ -5,19 +5,15 @@
 import torch
 from torch import nn
 from torch.utils._pytree import tree_flatten, tree_unflatten
-from torchvision.prototype import features
 from torchvision.prototype.transforms._utils import _isinstance
 from torchvision.utils import _log_api_usage_once
 
 
 class Transform(nn.Module):
 
     # Class attribute defining transformed types. Other types are passed-through without any transformation
-    _transformed_types: Tuple[Union[Type, Callable[[Any], bool]], ...] = (
-        features.is_simple_tensor,
-        features._Feature,
-        PIL.Image.Image,
-    )
+    # We support both Types and callables that are able to do further checks on the type of the input.
+    _transformed_types: Tuple[Union[Type, Callable[[Any], bool]], ...] = (torch.Tensor, PIL.Image.Image)
 
     def __init__(self) -> None:
         super().__init__()