add device for TensorDict in split_and_pad_trajectories (#138)

ClemensSchwarke · web-flow · commit e1e7071bea26 · 2025-11-26T12:15:21.000+01:00
diff --git a/rsl_rl/utils/utils.py b/rsl_rl/utils/utils.py
@@ -129,7 +129,7 @@ def split_and_pad_trajectories(
             # Remove the added trajectory
             padded_trajectories[k] = padded_trajectories[k][:, :-1]
         padded_trajectories = TensorDict(
-            padded_trajectories, batch_size=[tensor.batch_size[0], len(trajectory_lengths_list)]
+            padded_trajectories, batch_size=[tensor.batch_size[0], len(trajectory_lengths_list)], device=tensor.device
         )
     else:
         # Split the tensor into trajectories

Original file line number	Diff line number	Diff line change
`@@ -129,7 +129,7 @@ def split_and_pad_trajectories(`
`129`	`129`	`# Remove the added trajectory`
`130`	`130`	`padded_trajectories[k] = padded_trajectories[k][:, :-1]`
`131`	`131`	`padded_trajectories = TensorDict(`
`132`		`- padded_trajectories, batch_size=[tensor.batch_size[0], len(trajectory_lengths_list)]`
	`132`	`+ padded_trajectories, batch_size=[tensor.batch_size[0], len(trajectory_lengths_list)], device=tensor.device`
`133`	`133`	`)`
`134`	`134`	`else:`
`135`	`135`	`# Split the tensor into trajectories`