fix: only on first process, not once on every node

2024-10-01 01:06:10 -04:00 · 2023-04-05 02:36:22 +00:00 · 2023-04-05 02:36:22 +00:00 · 97d4499d79
commit 97d4499d79
parent d0402288bd
1 changed files with 1 additions and 1 deletions
--- a/train.py
+++ b/train.py
@ -137,7 +137,7 @@ def train(accelerator, config):


    # log gradients
-    if accelerator.is_local_main_process and config["wandb"]:
+    if accelerator.is_main_process and config["wandb"]:
        wandb.watch(model, log_freq=config["log_grads_every"])

    for epoch in range(config["num_epochs"]):