Merge 2fe78bbfac into 145fe4cd17

2025-04-03 10:31:22 -07:00 · 2025-04-03 10:31:22 -07:00 · 713eab7e3b
parent 145fe4cd17 2fe78bbfac
commit 713eab7e3b
1 changed files with 2 additions and 2 deletions
--- a/lerobot/common/policies/pi0/modeling_pi0.py
+++ b/lerobot/common/policies/pi0/modeling_pi0.py
@ -330,8 +330,8 @@ class PI0Policy(PreTrainedPolicy):

        # For backward pass
        loss = losses.mean()
-        # For logging
-        loss_dict["l2_loss"] = loss.item()
+        # For logging. Use detach so won't create scalar to break graph when using torch.compile
+        loss_dict["l2_loss"] = loss.detach()

        return loss, loss_dict