Merge f29e0d5bd4 into 768e36660d

2025-04-15 10:18:22 +08:00 · 2025-04-15 10:18:22 +08:00 · 959e8b0d7a
parent 768e36660d f29e0d5bd4
commit 959e8b0d7a
1 changed files with 2 additions and 1 deletions
--- a/lerobot/common/policies/pi0/modeling_pi0.py
+++ b/lerobot/common/policies/pi0/modeling_pi0.py
@ -325,7 +325,8 @@ class PI0Policy(PreTrainedPolicy):
            loss_dict["losses_after_in_ep_bound"] = losses.clone()

        # Remove padding
-        losses = losses[:, :, : self.config.max_action_dim]
+        original_action_dim = self.config.action_feature.shape[0]
+        losses = losses[:, :, :original_action_dim]
        loss_dict["losses_after_rm_padding"] = losses.clone()

        # For backward pass