admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-25 02:36:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

fixing wm_loss not propagating

Browse Source
This commit is contained in:
Maxime Ellerbach
2026-05-13 12:38:28 +00:00
committed by Maximellerbach
parent 6657be734c
commit 921b823fb4
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/vla_jepa/modeling_vla_jepa.py
+2 -2
View File
@@ -448,8 +448,8 @@ class VLAJEPAPolicy(PreTrainedPolicy):
logs["action_loss"] = native_output["action_loss"].detach().item() logs["action_loss"] = native_output["action_loss"].detach().item()
if "wm_loss" in native_output: if "wm_loss" in native_output:
wm_loss = native_output["wm_loss"] total_loss = total_loss + native_output["wm_loss"]
logs["wm_loss"] = wm_loss.detach().item() logs["wm_loss"] = native_output["wm_loss"].detach().item()
logs["loss"] = ( logs["loss"] = (
total_loss.detach().item() total_loss.detach().item()