Repository: Project-MONAI/MONAI

Issue:
UserWarning: TypedStorage is deprecated.
**Describe the bug**
see also https://github.com/Project-MONAI/MONAI/issues/5862
pytorch docker 23.01

MONAI version: 1.1.0+271.g8deedd61
Numpy version: 1.22.2
Pytorch version: 1.14.0a0+44dac51
MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False
MONAI rev id: 8deedd61e7bab8ae0fe58d24bb7e68db5d6fddb2
MONAI __file__: /workspace/MONAI/monai/__init__.py

```
root@4242496:/workspace/MONAI# python -W error -m tests.test_meta_tensor
...............................................................................Traceback (most recent call last):
  File "/usr/lib/python3.8/multiprocessing/queues.py", line 239, in _feed
    obj = _ForkingPickler.dumps(obj)
  File "/usr/lib/python3.8/multiprocessing/reduction.py", line 51, in dumps
    cls(buf, protocol).dump(obj)
  File "/workspace/MONAI/monai/data/__init__.py", line 132, in reduce_meta_tensor
    storage = meta_tensor.untyped_storage() if hasattr(meta_tensor, "untyped_storage") else meta_tensor.storage()
  File "/usr/local/lib/python3.8/dist-packages/torch/_tensor.py", line 226, in storage
    return handle_torch_function(Tensor.storage, (self,), self)
  File "/usr/local/lib/python3.8/dist-packages/torch/overrides.py", line 1537, in handle_torch_function
    result = torch_func_method(public_api, types, args, kwargs)
  File "/workspace/MONAI/monai/data/meta_tensor.py", line 276, in __torch_function__
    ret = super().__torch_function__(func, types, args, kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/_tensor.py", line 1283, in __torch_function__
    ret = func(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/_tensor.py", line 228, in storage
    torch.storage._warn_typed_storage_removal()
  File "/usr/local/lib/python3.8/dist-packages/torch/storage.py", line 315, in _warn_typed_storage_removal
    warnings.warn(message, UserWarning)
UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.
Traceback (most recent call last):
  File "/usr/lib/python3.8/multiprocessing/queues.py", line 239, in _feed
    obj = _ForkingPickler.dumps(obj)
  File "/usr/lib/python3.8/multiprocessing/reduction.py", line 51, in dumps
    cls(buf, protocol).dump(obj)
  File "/workspace/MONAI/monai/data/__init__.py", line 132, in reduce_meta_tensor
    storage = meta_tensor.untyped_storage() if hasattr(meta_tensor, "untyped_storage") else meta_tensor.storage()
  File "/usr/local/lib/python3.8/dist-packages/torch/_tensor.py", line 226, in storage
    return handle_torch_function(Tensor.storage, (self,), self)
  File "/usr/local/lib/python3.8/dist-packages/torch/overrides.py", line 1537, in handle_torch_function
    result = torch_func_method(public_api, types, args, kwargs)
  File "/workspace/MONAI/monai/data/meta_tensor.py", line 276, in __torch_function__
    ret = super().__torch_function__(func, types, args, kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/_tensor.py", line 1283, in __torch_function__
    ret = func(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/_tensor.py", line 228, in storage
    torch.storage._warn_typed_storage_removal()
  File "/usr/local/lib/python3.8/dist-packages/torch/storage.py", line 315, in _warn_typed_storage_removal
    warnings.warn(message, UserWarning)
UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.
```

```

```


Discussion:

Resolving patch:
diff --git a/monai/data/__init__.py b/monai/data/__init__.py
--- a/monai/data/__init__.py
+++ b/monai/data/__init__.py
@@ -122,18 +122,31 @@
     from multiprocessing.reduction import ForkingPickler
 
     def _rebuild_meta(cls, storage, dtype, metadata):
-        storage_offset, size, stride, meta_dict = metadata
+        storage_offset, size, stride, requires_grad, meta_dict = metadata
+        storage = storage._untyped_storage if hasattr(storage, "_untyped_storage") else storage
         t = cls([], dtype=dtype, device=storage.device)
-        t.set_(storage._untyped() if hasattr(storage, "_untyped") else storage, storage_offset, size, stride)
+        t.set_(storage, storage_offset, size, stride)
+        t.requires_grad = requires_grad
         t.__dict__ = meta_dict
         return t
 
     def reduce_meta_tensor(meta_tensor):
-        storage = meta_tensor.untyped_storage() if hasattr(meta_tensor, "untyped_storage") else meta_tensor.storage()
+        if hasattr(meta_tensor, "untyped_storage"):
+            storage = meta_tensor.untyped_storage()
+        elif hasattr(meta_tensor, "_typed_storage"):  # gh pytorch 44dac51/torch/_tensor.py#L231-L233
+            storage = meta_tensor._typed_storage()
+        else:
+            storage = meta_tensor.storage()
         dtype = meta_tensor.dtype
-        if storage.is_cuda:
+        if meta_tensor.is_cuda:
             raise NotImplementedError("sharing CUDA metatensor across processes not implemented")
-        metadata = (meta_tensor.storage_offset(), meta_tensor.size(), meta_tensor.stride(), meta_tensor.__dict__)
+        metadata = (
+            meta_tensor.storage_offset(),
+            meta_tensor.size(),
+            meta_tensor.stride(),
+            meta_tensor.requires_grad,
+            meta_tensor.__dict__,
+        )
         return _rebuild_meta, (type(meta_tensor), storage, dtype, metadata)
 
     ForkingPickler.register(MetaTensor, reduce_meta_tensor)


Test patch:
diff --git a/tests/test_meta_tensor.py b/tests/test_meta_tensor.py
--- a/tests/test_meta_tensor.py
+++ b/tests/test_meta_tensor.py
@@ -247,7 +247,7 @@ def test_torchscript(self, device):
                     "your pytorch version if this is important to you."
                 )
                 im_conv = im_conv.as_tensor()
-            self.check(out, im_conv, ids=False)
+        self.check(out, im_conv, ids=False)
 
     def test_pickling(self):
         m, _ = self.get_im()
@@ -258,7 +258,7 @@ def test_pickling(self):
             if not isinstance(m2, MetaTensor) and not pytorch_after(1, 8, 1):
                 warnings.warn("Old version of pytorch. pickling converts `MetaTensor` to `torch.Tensor`.")
                 m = m.as_tensor()
-            self.check(m2, m, ids=False)
+        self.check(m2, m, ids=False)
 
     @skip_if_no_cuda
     def test_amp(self):