How to log Molecules via wandb in case of multi-gpu training?

[glukhove]

Hello!
I'd like to log Molecule on each validation step.
So for 1-GPU training I was able to use wandb_logger created outside of the scope. So, instead of self.log I used wandb_logger.log_table(...) and it was fine.
But when I use multiple GPUs wandb_logger would work only on rank 0 gpu, all other will not upload anything.
I was trying to create a workaround, but wasn't able to find anything decent.