Fix SmolVLA loss not sent to wandb (#1198)

This commit is contained in:
Ben Zhang
2025-06-05 02:13:03 -07:00
committed by GitHub
parent bfd26eef5a
commit b536f47e3f
2 changed files with 2 additions and 2 deletions

View File

@@ -338,7 +338,7 @@ class SmolVLAPolicy(PreTrainedPolicy):
# For backward pass # For backward pass
loss = losses.mean() loss = losses.mean()
# For backward pass # For backward pass
loss_dict["loss"] = loss loss_dict["loss"] = loss.item()
return loss, loss_dict return loss, loss_dict
def prepare_images(self, batch): def prepare_images(self, batch):

View File

@@ -115,7 +115,7 @@ class WandBLogger:
for k, v in d.items(): for k, v in d.items():
if not isinstance(v, (int, float, str)): if not isinstance(v, (int, float, str)):
logging.warning( logging.warning(
f'WandB logging of key "{k}" was ignored as its type is not handled by this wrapper.' f'WandB logging of key "{k}" was ignored as its type "{type(v)}" is not handled by this wrapper.'
) )
continue continue
self._wandb.log({f"{mode}/{k}": v}, step=step) self._wandb.log({f"{mode}/{k}": v}, step=step)