make wandb histograms work again

This commit is contained in:
James Betker 2022-05-30 16:44:39 -06:00
parent 29b55d42a5
commit 479c3195f3

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: