myLightningOPD / slime /utils /tracking_utils.py
ayh015's picture
Upload folder using huggingface_hub
6011e08 verified
Raw
History Blame Contribute Delete
798 Bytes
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import wandb
from slime.utils.tensorboard_utils import _TensorboardAdapter
from . import wandb_utils
def init_tracking(args, primary: bool = True, **kwargs):
if primary:
wandb_utils.init_wandb_primary(args, **kwargs)
else:
wandb_utils.init_wandb_secondary(args, **kwargs)
# TODO further refactor, e.g. put TensorBoard init to the "init" part
def log(args, metrics, step_key: str):
if args.use_wandb:
wandb.log(metrics)
if args.use_tensorboard:
metrics_except_step = {k: v for k, v in metrics.items() if k != step_key}
_TensorboardAdapter(args).log(data=metrics_except_step, step=metrics[step_key])