fix tensorboard error (jpata#254)

* fix tensorboard error * suppress warning * remove enviornment variable * Update pyg_pipeline.py --------- Co-authored-by: Joosep Pata <[email protected]>
farakiko · Oct 25, 2023 · 38d2e59 · 38d2e59
1 parent 370f47f
commit 38d2e59
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/mlpf/pyg/training.py b/mlpf/pyg/training.py
@@ -370,10 +370,10 @@ def train_mlpf(rank, world_size, model, optimizer, train_loader, valid_loader, n
         #     with record_function("model_train"):
         # print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
 
-        for k, v in losses_t.items():
-            tensorboard_writer.add_scalar(f"epoch/train_loss_rank_{rank}_" + k, v, epoch)
-
         if (rank == 0) or (rank == "cpu"):
+            for k, v in losses_t.items():
+                tensorboard_writer.add_scalar("epoch/train_loss_" + k, v, epoch)
+
             for loss in losses_of_interest:
                 losses["train"][loss].append(losses_t[loss])
                 losses["valid"][loss].append(losses_v[loss])