fix: fix storage issue in torchtensor class (#1833)

ai-naymul · web-flow · commit 3cfa0b8ff877 · 2023-12-09T18:14:22.000+01:00
Signed-off-by: Naymul Islam &lt;naymul504@gmail.com&gt;
diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py
@@ -293,6 +293,16 @@ def __torch_function__(cls, func, types, args=(), kwargs=None):
         )
         return super().__torch_function__(func, types_, args, kwargs)
 
+    def __deepcopy__(self, memo):
+        """
+        Custom implementation of deepcopy for TorchTensor to avoid storage sharing issues.
+        """
+        # Create a new tensor with the same data and properties
+        new_tensor = self.clone()
+        # Set the class to the custom TorchTensor class
+        new_tensor.__class__ = self.__class__
+        return new_tensor
+
     @classmethod
     def _docarray_from_ndarray(cls: Type[T], value: np.ndarray) -> T:
         """Create a `tensor from a numpy array
diff --git a/tests/integrations/typing/test_torch_tensor.py b/tests/integrations/typing/test_torch_tensor.py
@@ -1,4 +1,6 @@
 import torch
+from docarray.typing.tensor.torch_tensor import TorchTensor
+import copy
 
 from docarray import BaseDoc
 from docarray.typing import TorchEmbedding, TorchTensor
@@ -25,3 +27,19 @@ class MyDocument(BaseDoc):
     assert isinstance(d.embedding, TorchEmbedding)
     assert isinstance(d.embedding, torch.Tensor)
     assert (d.embedding == torch.zeros((128,))).all()
+
+
+def test_torchtensor_deepcopy():
+    # Setup
+    original_tensor_float = TorchTensor(torch.rand(10))
+    original_tensor_int = TorchTensor(torch.randint(0, 100, (10,)))
+
+    # Exercise
+    copied_tensor_float = copy.deepcopy(original_tensor_float)
+    copied_tensor_int = copy.deepcopy(original_tensor_int)
+
+    # Verify
+    assert torch.equal(original_tensor_float, copied_tensor_float)
+    assert original_tensor_float is not copied_tensor_float
+    assert torch.equal(original_tensor_int, copied_tensor_int)
+    assert original_tensor_int is not copied_tensor_int
