[BugFix] Fix non-blocking arg in copy_ (#590)

tensordict/base.py

@@ -2126,8 +2126,14 @@ def _create_nested_tuple(self, key):
         if len(key) > 1:
             td._create_nested_tuple(key[1:])
 
-    def copy_(self, tensordict: T) -> T:
-        """See :obj:`TensorDictBase.update_`."""
+    def copy_(self, tensordict: T, non_blocking: bool = None) -> T:
+        """See :obj:`TensorDictBase.update_`.
+
+        The non-blocking argument will be ignored and is just present for
+        compatibility with :func:`torch.Tensor.copy_`.
+        """
+        if non_blocking is False:
+            raise ValueError("non_blocking=False isn't supported in TensorDict.")
         return self.update_(tensordict)
 
     def copy_at_(self, tensordict: T, idx: IndexType) -> T:

tensordict/memmap_deprec.py

@@ -526,7 +526,9 @@ def cuda(self) -> torch.Tensor:
     def numpy(self) -> np.ndarray:
         return self._tensor.numpy()
 
-    def copy_(self, other: torch.Tensor | MemmapTensor) -> MemmapTensor:
+    def copy_(
+        self, other: torch.Tensor | MemmapTensor, non_blocking: bool = False
+    ) -> MemmapTensor:
         if isinstance(other, MemmapTensor) and other.filename == self.filename:
             if not self.shape == other.shape:
                 raise ValueError(