training the model

2025-05-28 14:24:30 +02:00
parent 8d6acb3a4f
commit 4719f41a43
2 changed files with 59 additions and 1 deletions
--- a/lerobot/common/policies/smolvla/modeling_smolvla.py
+++ b/lerobot/common/policies/smolvla/modeling_smolvla.py
@@ -335,7 +335,7 @@ class SmolVLAPolicy(PreTrainedPolicy):
        loss = losses.mean()
        # For backward pass
        loss_dict["loss"] = loss
-        return loss_dict
+        return loss, loss_dict

    def prepare_images(self, batch):
        """Apply SmolVLA preprocessing to the images, like resizing to 224x224 and padding to keep aspect ratio, and