docarray · alaeddine-13 · Apr 5, 2022 · Mar 10, 2022 · Mar 10, 2022 · Mar 10, 2022
diff --git a/docarray/array/mixins/setitem.py b/docarray/array/mixins/setitem.py
@@ -120,6 +120,7 @@ def __setitem__(
             raise IndexError(f'Unsupported index type {typename(index)}: {index}')
 
     def _set_by_pair(self, idx1, idx2, value):
+
         if isinstance(idx1, str) and not idx1.startswith('@'):
             # second is an ID
             # allows da[id1, id2] = [d1, d2]
@@ -136,6 +137,7 @@ def _set_by_pair(self, idx1, idx2, value):
                 and all(isinstance(attr, str) for attr in idx2)
                 and all(hasattr(self[idx1], attr) for attr in idx2)
             ):
+
                 for attr, _v in zip(idx2, value):
                     self._set_doc_attr_by_id(idx1, attr, _v)
             else:

diff --git a/docarray/document/__init__.py b/docarray/document/__init__.py
@@ -3,6 +3,7 @@
 from .data import DocumentData
 from .mixins import AllMixins
 from ..base import BaseDCType
+from ..math.ndarray import detach_tensor_if_present
 
 if TYPE_CHECKING:
     from ..types import ArrayType, StructValueType, DocumentContentType
@@ -88,3 +89,16 @@ def __init__(
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+
+        for attribute in ['embedding', 'tensor']:
+            if hasattr(self, attribute):
+                setattr(
+                    state['_data'],
+                    attribute,
+                    detach_tensor_if_present(getattr(state['_data'], attribute)),
+                )
+
+        return state
diff --git a/docarray/math/ndarray.py b/docarray/math/ndarray.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Tuple, Sequence, Optional, List
+from typing import TYPE_CHECKING, Tuple, Sequence, Optional, List, Any
 
 import numpy as np
 
@@ -81,7 +81,9 @@ def ravel(value: 'ArrayType', docs: 'DocumentArray', field: str) -> None:
             docs[d.id, field] = value[j, ...]
 
 
-def get_array_type(array: 'ArrayType') -> Tuple[str, bool]:
+def get_array_type(
+    array: 'ArrayType', raise_error_if_not_array: bool = True
+) -> Tuple[str, bool]:
     """Get the type of ndarray without importing the framework
 
     :param array: any array, scipy, numpy, tf, torch, etc.
@@ -121,7 +123,10 @@ def get_array_type(array: 'ArrayType') -> Tuple[str, bool]:
     if 'scipy' in module_tags and 'sparse' in module_tags:
         return 'scipy', True
 
-    raise TypeError(f'can not determine the array type: {module_tags}.{class_name}')
+    if raise_error_if_not_array:
+        raise TypeError(f'can not determine the array type: {module_tags}.{class_name}')
+    else:
+        return 'python', False
 
 
 def to_numpy_array(value) -> 'np.ndarray':
@@ -258,3 +263,16 @@ def check_arraylike_equality(x: 'ArrayType', y: 'ArrayType'):
         return same_array
     else:
         return same_array
+
+
+def detach_tensor_if_present(x: Any) -> Any:
+    """Check if input is a dense torch array and detaches the tensor from the current graph.
+    :param array: input array
+    :return: (num_rows, ndim)
+    """
+    x_type, x_sparse = get_array_type(x, raise_error_if_not_array=False)
+    if x_type == 'torch' and x_sparse == False:
+        import torch
+
+        x = torch.tensor(x.detach().numpy())
+    return x