Update on "add memory metrics to TensorBoard"

<img width="1391" alt="Screenshot 2024-02-15 at 5 19 09 PM" src="https://github.com/pytorch-labs/torchtrain/assets/150487191/af8a2efb-13ff-4e8f-84f2-b245784747ed"> [ghstack-poisoned]
pytorch · Feb 17, 2024 · b77c89f · b77c89f
1 parent e1a577d
commit b77c89f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -222,8 +222,8 @@ def main(args):
                 gpu_mem_stats = gpu_metrics.get_current_stats(return_data=True)
 
                 metrics = {
-                    "loss/global_avg": global_avg_loss,
-                    "loss/global_max": global_max_loss,
+                    "loss_metrics/global_avg_loss": global_avg_loss,
+                    "loss_metrics/global_max_loss": global_max_loss,
                     "wps": wps,
                     "memory_current/active(%)": gpu_mem_stats.active_curr,
                     "memory_current/allocated(%)": gpu_mem_stats.allocated_curr,