From d847949fe96dfc78e6f59d9748ce32a35478475f Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Sat, 6 May 2023 10:15:48 -0700 Subject: [PATCH] fixing wandb envs (#1880) * adding back support for wandb envs * remove .init --- nerfstudio/utils/writer.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index 2036e4659..6deee6900 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -18,6 +18,7 @@ Generic Writer class from __future__ import annotations import enum +import os from abc import abstractmethod from pathlib import Path from time import time @@ -287,7 +288,12 @@ class WandbWriter(Writer): """WandDB Writer Class""" def __init__(self, log_dir: Path, experiment_name: str, project_name: str = "nerfstudio-project"): - wandb.init(project=project_name, dir=str(log_dir), reinit=True, name=experiment_name) + wandb.init( + project=os.environ.get("WANDB_PROJECT", project_name), + dir=os.environ.get("WANDB_DIR", str(log_dir)), + name=os.environ.get("WANDB_NAME", experiment_name), + reinit=True, + ) def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: image = torch.permute(image, (2, 0, 1))