diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py index a78781f6a9b..32d1cbff7be 100644 --- a/docarray/typing/tensor/torch_tensor.py +++ b/docarray/typing/tensor/torch_tensor.py @@ -271,3 +271,10 @@ def _docarray_from_ndarray(cls: Type[T], value: np.ndarray) -> T: def _docarray_to_ndarray(self) -> np.ndarray: """cast itself to a numpy array""" return self.detach().cpu().numpy() + + def new_empty(self, *args, **kwargs): + """ + This method enables the deepcopy of `TorchTensor` by returning another instance of this subclass. + If this function is not implemented, the deepcopy will throw an RuntimeError from Torch. + """ + return self.__class__(*args, **kwargs) diff --git a/tests/units/typing/tensor/test_torch_tensor.py b/tests/units/typing/tensor/test_torch_tensor.py index 0f3c9882e2a..02605142113 100644 --- a/tests/units/typing/tensor/test_torch_tensor.py +++ b/tests/units/typing/tensor/test_torch_tensor.py @@ -187,6 +187,22 @@ class MMdoc(BaseDoc): assert not (doc.embedding == doc_copy.embedding).all() +def test_deepcopy_tensor(): + from docarray import BaseDoc + + class MMdoc(BaseDoc): + embedding: TorchTensor + + doc = MMdoc(embedding=torch.randn(32)) + doc_copy = doc.copy(deep=True) + + assert doc.embedding.data_ptr() != doc_copy.embedding.data_ptr() + assert (doc.embedding == doc_copy.embedding).all() + + doc_copy.embedding = torch.randn(32) + assert not (doc.embedding == doc_copy.embedding).all() + + @pytest.mark.parametrize('requires_grad', [True, False]) def test_json_serialization(requires_grad): orig_doc = MyDoc(tens=torch.rand(10, requires_grad=requires_grad))