From 47cc89c94d7e123da5746a09ed78807f9b5403f4 Mon Sep 17 00:00:00 2001
From: Pqlet <kartinkudi@yandex.ru>
Date: Sun, 25 May 2025 11:14:34 +0300
Subject: [PATCH] Add step logging

---
 dictionary_learning/training.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/dictionary_learning/training.py b/dictionary_learning/training.py
index 0671f31..8c212dc 100644
--- a/dictionary_learning/training.py
+++ b/dictionary_learning/training.py
@@ -26,7 +26,7 @@ def new_wandb_process(config, log_queue, entity, project):
             log = log_queue.get(timeout=1)
             if log == "DONE":
                 break
-            wandb.log(log)
+            wandb.log(**log)
         except Empty:
             continue
     wandb.finish()
@@ -78,7 +78,7 @@ def log_stats(
                 log[f"{name}"] = value
 
             if log_queues:
-                log_queues[i].put(log)
+                log_queues[i].put({"data":log, "step":step})
 
 def get_norm_factor(data, steps: int) -> float:
     """Per Section 3.1, find a fixed scalar factor so activation vectors have unit mean squared norm.