[BugFix] LLMCollector trajectory collection methods only returns 1 trajectory when multiple are completed at the same step (#3018)

Author: albertbou92

torchrl/collectors/llm/base.py

@@ -347,7 +347,7 @@ def _rollout_yield_trajs(self) -> TensorDictBase:  # A simplified version of rol
                 queue.append(_data)
                 dones[i] = _data["next", "done"].any()
             if dones.any():
-                for idx in dones.nonzero()[0].tolist():
+                for idx in dones.nonzero(as_tuple=True)[0].tolist():
                     if not self.yield_only_last_steps:
                         self._trajectory_queue.append(
                             lazy_stack(self._yield_queues[idx], -1)
@@ -401,7 +401,7 @@ def _rollout_yield_trajs_async(
                 queue.append(_data)
                 dones[i] = _data["next", "done"].any()
             if dones.any():
-                for idx in dones.nonzero()[0].tolist():
+                for idx in dones.nonzero(as_tuple=True)[0].tolist():
                     if not self.yield_only_last_steps:
                         self._trajectory_queue.append(
                             lazy_stack(self._yield_queues[idx], -1)