From efe361911e9e2af9bada50b63ddb88135a53175f Mon Sep 17 00:00:00 2001
From: alexcbb <alexchapin@hotmail.fr>
Date: Mon, 24 Jul 2023 16:02:47 +0200
Subject: [PATCH] Deactivate wandb
---
train_sa.py | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/train_sa.py b/train_sa.py
index 17886ac..269319c 100644
--- a/train_sa.py
+++ b/train_sa.py
@@ -60,7 +60,6 @@ def main():
#### Create model
model = LitSlotAttentionAutoEncoder(resolution, num_slots, num_iterations, cfg=cfg)
- wandb_logger = WandbLogger(project="slot-attention")
checkpoint_callback = ModelCheckpoint(
save_top_k=10,
@@ -71,7 +70,7 @@ def main():
)
trainer = pl.Trainer(accelerator="gpu", devices=num_gpus, profiler="simple",
- default_root_dir="./logs", logger=wandb_logger,
+ default_root_dir="./logs", logger=WandbLogger(project="slot-attention") if args.wandb else None,
strategy="ddp" if num_gpus > 1 else "default", callbacks=[checkpoint_callback], deterministic=True,
log_every_n_steps=100, max_steps=num_train_steps)
--
GitLab