fix consolidated locks

Author: vmoens

test/test_rb.py

@@ -1791,6 +1791,10 @@ def test_batch_errors():
 
 
 def test_add_warning():
+    from torchrl._utils import RL_WARNINGS
+
+    if not RL_WARNINGS:
+        return
     rb = ReplayBuffer(storage=ListStorage(10), batch_size=3)
     with pytest.warns(
         UserWarning,

torchrl/data/llm/history.py

@@ -1219,15 +1219,6 @@ def append(
                 f"The new history to append must have one less dimension than self. Got self.ndim={self.ndim} and history.ndim={history.ndim}."
             )
         dim = _maybe_correct_neg_dim(dim, self.batch_size)
-        # if self.ndim > 1 and dim >= self.ndim - 1:
-        #     # then we need to append each element independently
-        #     result = []
-        #     for hist, new_hist in zip(self.unbind(0), history.unbind(0)):
-        #         hist_c = hist.append(new_hist, inplace=inplace, dim=dim - 1)
-        #         result.append(hist_c)
-        #     if inplace:
-        #         return self
-        #     return lazy_stack(result)
         if inplace:
             if (
                 isinstance(self._tensordict, LazyStackedTensorDict)

torchrl/envs/batched_envs.py

@@ -1843,7 +1843,7 @@ def _step_no_buffers(
         if self.consolidate:
             try:
                 data = tensordict.consolidate(
-                    share_memory=True, inplace=True, num_threads=1
+                    share_memory=True, inplace=False, num_threads=1
                 )
             except Exception as err:
                 raise RuntimeError(_CONSOLIDATE_ERR_CAPTURE) from err
@@ -2677,6 +2677,7 @@ def _run_worker_pipe_direct(
             # data = data[idx]
             data, reset_kwargs = data
             if data is not None:
+                data.unlock_()
                 data._fast_apply(
                     lambda x: x.clone() if x.device.type == "cuda" else x, out=data
                 )