mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-17 17:50:09 +00:00
Convert tensors to scalars (#2903)
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
This commit is contained in:
@@ -378,16 +378,16 @@ class SmolVLAPolicy(PreTrainedPolicy):
|
|||||||
actions_is_pad = batch.get("actions_id_pad")
|
actions_is_pad = batch.get("actions_id_pad")
|
||||||
loss_dict = {}
|
loss_dict = {}
|
||||||
losses = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions, noise, time)
|
losses = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions, noise, time)
|
||||||
loss_dict["losses_after_forward"] = losses.clone()
|
loss_dict["losses_after_forward"] = losses.clone().mean().item()
|
||||||
|
|
||||||
if actions_is_pad is not None:
|
if actions_is_pad is not None:
|
||||||
in_episode_bound = ~actions_is_pad
|
in_episode_bound = ~actions_is_pad
|
||||||
losses = losses * in_episode_bound.unsqueeze(-1)
|
losses = losses * in_episode_bound.unsqueeze(-1)
|
||||||
loss_dict["losses_after_in_ep_bound"] = losses.clone()
|
loss_dict["losses_after_in_ep_bound"] = losses.clone().mean().item()
|
||||||
|
|
||||||
# Remove padding
|
# Remove padding
|
||||||
losses = losses[:, :, : self.config.max_action_dim]
|
losses = losses[:, :, : self.config.max_action_dim]
|
||||||
loss_dict["losses_after_rm_padding"] = losses.clone()
|
loss_dict["losses_after_rm_padding"] = losses.clone().mean().item()
|
||||||
|
|
||||||
if reduction == "none":
|
if reduction == "none":
|
||||||
# Return per-sample losses (B,) by averaging over time and action dims
|
# Return per-sample losses (B,) by averaging over time and action dims
|
||||||
|
|||||||
Reference in New Issue
Block a user