diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py index 7ad743721a..7ed3bd3800 100644 --- 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 index 2a84489cd9..0e485fcd07 100644 --- 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