Use _tensor ref

vfdev-5 · vfdev-5 · commit 38f8e2124283 · 2022-10-14T15:40:24.000Z
diff --git a/torchvision/prototype/features/_bounding_box.py b/torchvision/prototype/features/_bounding_box.py
@@ -24,6 +24,7 @@ def _wrap(cls, tensor: torch.Tensor, *, format: BoundingBoxFormat, spatial_size:
         bounding_box = tensor.as_subclass(cls)
         bounding_box.format = format
         bounding_box.spatial_size = spatial_size
+        bounding_box._tensor = tensor
         return bounding_box
 
     def __new__(
diff --git a/torchvision/prototype/features/_encoded.py b/torchvision/prototype/features/_encoded.py
@@ -16,7 +16,9 @@
 class EncodedData(_Feature):
     @classmethod
     def _wrap(cls: Type[D], tensor: torch.Tensor) -> D:
-        return tensor.as_subclass(cls)
+        output = tensor.as_subclass(cls)
+        output._tensor = tensor
+        return output
 
     def __new__(
         cls,
diff --git a/torchvision/prototype/features/_feature.py b/torchvision/prototype/features/_feature.py
@@ -44,6 +44,7 @@ def __new__(
     ) -> _Feature:
         tensor = cls._to_tensor(data, dtype=dtype, device=device, requires_grad=requires_grad)
         output = tensor.as_subclass(_Feature)
+        output._tensor = tensor
         return output
 
     @classmethod
@@ -108,7 +109,9 @@ def __torch_function__(
             # Inplace `func`'s, canonically identified with a trailing underscore in their name like `.add_(...)`,
             # will retain the input type. Thus, we need to unwrap here.
             if isinstance(output, cls):
-                return output.as_subclass(torch.Tensor)
+                tensor = output.as_subclass(torch.Tensor)
+                output._tensor = tensor
+                return tensor
 
             return output
 
@@ -134,23 +137,19 @@ def _F(self) -> ModuleType:
     # this way we return the result without passing into __torch_function__
     @property
     def shape(self) -> _size:  # type: ignore[override]
-        with DisableTorchFunction():
-            return super().shape
+        return self._tensor.shape
 
     @property
     def ndim(self) -> int:  # type: ignore[override]
-        with DisableTorchFunction():
-            return super().ndim
+        return self._tensor.ndim
 
     @property
     def device(self, *args: Any, **kwargs: Any) -> _device:  # type: ignore[override]
-        with DisableTorchFunction():
-            return super().device
+        return self._tensor.device
 
     @property
     def dtype(self) -> _dtype:  # type: ignore[override]
-        with DisableTorchFunction():
-            return super().dtype
+        return self._tensor.dtype
 
     def horizontal_flip(self) -> _Feature:
         return self
diff --git a/torchvision/prototype/features/_image.py b/torchvision/prototype/features/_image.py
@@ -64,6 +64,7 @@ class Image(_Feature):
     def _wrap(cls, tensor: torch.Tensor, *, color_space: ColorSpace) -> Image:
         image = tensor.as_subclass(cls)
         image.color_space = color_space
+        image._tensor = tensor
         return image
 
     def __new__(
diff --git a/torchvision/prototype/features/_label.py b/torchvision/prototype/features/_label.py
@@ -18,6 +18,7 @@ class _LabelBase(_Feature):
     def _wrap(cls: Type[L], tensor: torch.Tensor, *, categories: Optional[Sequence[str]]) -> L:
         label_base = tensor.as_subclass(cls)
         label_base.categories = categories
+        label_base._tensor = tensor
         return label_base
 
     def __new__(
diff --git a/torchvision/prototype/features/_mask.py b/torchvision/prototype/features/_mask.py
@@ -11,7 +11,9 @@
 class Mask(_Feature):
     @classmethod
     def _wrap(cls, tensor: torch.Tensor) -> Mask:
-        return tensor.as_subclass(cls)
+        output = tensor.as_subclass(cls)
+        output._tensor = tensor
+        return output
 
     def __new__(
         cls,
diff --git a/torchvision/prototype/features/_video.py b/torchvision/prototype/features/_video.py
@@ -17,6 +17,7 @@ class Video(_Feature):
     def _wrap(cls, tensor: torch.Tensor, *, color_space: ColorSpace) -> Video:
         video = tensor.as_subclass(cls)
         video.color_space = color_space
+        video._tensor = tensor
         return video
 
     def __new__(
