deviandice/DL-Art-School
1
0
Fork 0
forked from mrq/DL-Art-School
Code Issues Pull Requests Packages Projects Releases Wiki Activity

make wandb histograms work again

Browse Source
This commit is contained in:
James Betker 2022-05-30 16:44:39 -06:00
parent 29b55d42a5
commit 479c3195f3
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
codes/train.py
View File

@ -238,6 +238,12 @@ class Trainer:
self.tb_logger.add_scalar(k, v, self.current_step)
if opt['wandb'] and self.rank <= 0:
import wandb
wandb_logs = {}
for k, v in logs.items():
if 'histogram' in k:
wandb_logs[k] = wandb.Histogram(v)
else:
wandb_logs[k] = v
if opt_get(opt, ['wandb_progress_use_raw_steps'], False):
wandb.log(logs, step=self.current_step)
else: