How to fix WandB plotting issues?

[djberenberg]

I'm having some trouble with a few of the visualizations available through the WandBLogger, I'm wondering if anyone has experienced similar issues.

The steps to reproduce this involve something like:

# previously logged in via `wandb login`
pl_module = SomeLightningModule()
wandb_logger = pl.loggers.WandBLogger()
wandb_logger.watch(pl_module.component1, 'gradients' , log_freq=1)
wandb_logger.watch(pl_module.component2, 'gradients' , log_freq=1) # could be a different log_freq
...
trainer = pl.Trainer(max_epochs=10, ...)
train, val = GET_DATA(...)
trainer.fit(pl_module, train, val)

1. The values in the metrics plots do not match the values reported to stdout.
   Below I am showing one example.
   Here is the WandB visualization, note the loss for epoch 5:

and the terminal output for epoch 5:

2. For log_freq=1, the gradient visualizations collapse. Screenshot seen below. Could this have to do with the fact that no gradients are reported during validation steps?

Other info:

• I am using multiple dataloaders under the min_size iteration paradigm, one (validation) dataset is much longer than the other.
• At the end of my training/validation steps I call self.log_dict in my LightningModule to log/aggregate metrics