log to wandb

Signed-off-by: Ye Yu <yeyu@nvidia.com>

Author: yeyu-nvidia

examples/speculative_decoding/eagle_utils.py

@@ -211,6 +211,9 @@ def on_log(self, args, state, control, **kwargs):
 
         # log to wandb
         if wandb and is_master():
+            logs = kwargs.get("logs") or {}
+            if logs:
+                wandb.log({k: v for k, v in logs.items() if v is not None}, step=state.global_step)
             for i, draft_acc in enumerate(average_acc):
                 for j, step_acc in enumerate(draft_acc):
                     wandb.log(