From b86a2c0b470ddb0579a5beec61ed30aee0de20a5 Mon Sep 17 00:00:00 2001
From: Simon Alibert <75076266+aliberts@users.noreply.github.com>
Date: Fri, 14 Feb 2025 18:00:12 +0100
Subject: [PATCH] Fix wandb logging (#730)

---
 lerobot/scripts/train.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/lerobot/scripts/train.py b/lerobot/scripts/train.py
index 7a31f2f5..f3c57fe2 100644
--- a/lerobot/scripts/train.py
+++ b/lerobot/scripts/train.py
@@ -232,7 +232,9 @@ def train(cfg: TrainPipelineConfig):
         if is_log_step:
             logging.info(train_tracker)
             if wandb_logger:
-                wandb_log_dict = {**train_tracker.to_dict(), **output_dict}
+                wandb_log_dict = train_tracker.to_dict()
+                if output_dict:
+                    wandb_log_dict.update(output_dict)
                 wandb_logger.log_dict(wandb_log_dict, step)
             train_tracker.reset_averages()