[Performance] Avoid cloning trajs in SliceSampler

vmoens · vmoens · commit 4fd54fef493d · 2024-12-20T10:27:36.000Z
ghstack-source-id: 2e133fc Pull Request resolved: #2671
diff --git a/torchrl/data/replay_buffers/samplers.py b/torchrl/data/replay_buffers/samplers.py
@@ -1243,7 +1243,7 @@ def _get_stop_and_length(self, storage, fallback=True):
                             "Could not get a tensordict out of the storage, which is required for SliceSampler to compute the trajectories."
                         )
                 vals = self._find_start_stop_traj(
-                    trajectory=trajectory.clone(),
+                    trajectory=trajectory,
                     at_capacity=storage._is_full,
                     cursor=getattr(storage, "_last_cursor", None),
                 )

Original file line number	Diff line number	Diff line change
`@@ -1243,7 +1243,7 @@ def _get_stop_and_length(self, storage, fallback=True):`
`1243`	`1243`	`"Could not get a tensordict out of the storage, which is required for SliceSampler to compute the trajectories."`
`1244`	`1244`	`)`
`1245`	`1245`	`vals = self._find_start_stop_traj(`
`1246`		`- trajectory=trajectory.clone(),`
	`1246`	`+ trajectory=trajectory,`
`1247`	`1247`	`at_capacity=storage._is_full,`
`1248`	`1248`	`cursor=getattr(storage, "_last_cursor", None),`
`1249`	`1249`	`)`