Unverified Commit eba418ac authored by Tristan Deleu's avatar Tristan Deleu Committed by GitHub
Browse files

Commit the last step on world_process_zero in WandbCallback (#9805)

* Commit the last step on world_process_zero in WandbCallback

* Use the environment variable WANDB_LOG_MODEL as a default value in WandbCallback
parent 8edc98bb
......@@ -516,6 +516,8 @@ class WandbCallback(TrainerCallback):
else:
self._wandb = wandb
self._initialized = False
# log outputs
self._log_model = os.getenv("WANDB_LOG_MODEL", "FALSE").upper() in ENV_VARS_TRUE_VALUES.union({"TRUE"})
def setup(self, args, state, model, reinit, **kwargs):
"""
......@@ -569,9 +571,6 @@ class WandbCallback(TrainerCallback):
model, log=os.getenv("WANDB_WATCH", "gradients"), log_freq=max(100, args.logging_steps)
)
# log outputs
self._log_model = os.getenv("WANDB_LOG_MODEL", "FALSE").upper() in ENV_VARS_TRUE_VALUES.union({"TRUE"})
def on_train_begin(self, args, state, control, model=None, **kwargs):
if self._wandb is None:
return
......@@ -583,7 +582,8 @@ class WandbCallback(TrainerCallback):
if self._wandb is None:
return
# commit last step
self._wandb.log({})
if state.is_world_process_zero:
self._wandb.log({})
if self._log_model and self._initialized and state.is_world_process_zero:
from .trainer import Trainer
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment