modified loss.item to loss.detach, so that no scalar will be created and when using torch.compile, it will not break graph

2025-02-14 08:43:09 -05:00 · 2025-02-14 08:43:09 -05:00 · f15dd56d0e
parent c574eb4984
commit f15dd56d0e
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