Make LoggingProxyWrapper report inf_step if requested (#194)

EmbarkStudios · Dec 7, 2023 · 24373cb · 24373cb
1 parent 349291f
commit 24373cb
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/emote/memory/memory.py b/emote/memory/memory.py
@@ -312,6 +312,8 @@ def get_report(
                 out_lists[key] = list(window)
                 out[key] = sum(window) / len(window)
                 out[f"{key}/cumulative"] = self.windowed_scalar_cumulative[key]
+            elif key == "inf_step":
+                out[key] = self.completed_inferences
 
         return out, out_lists