From 4a0259915af114dba9156e1ca05c0fb3983fb88c Mon Sep 17 00:00:00 2001 From: ClemensSchwarke Date: Thu, 13 Nov 2025 09:23:50 +0100 Subject: [PATCH] add device --- rsl_rl/utils/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rsl_rl/utils/utils.py b/rsl_rl/utils/utils.py index 7a044e83..dc4eeee2 100644 --- a/rsl_rl/utils/utils.py +++ b/rsl_rl/utils/utils.py @@ -116,7 +116,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