Error when trying to log wandb confusion matrix

I’m trying to figure out how to log a wandb confusion matrix as described here inside Lightning. I’m calling self.log inside of my lightning module’s validation_step and logging a custom metric which returns a wandb CustomChart from its compute method and getting the following error:

File "/home/ryan/anaconda3/envs/ml/lib/python3.9/site-packages/pytorch_lightning/core/", line 500, in reduce_on_epoch_end
    result[k] = torch.tensor(result[k]).float()
RuntimeError: Could not infer dtype of CustomChart

This is one of several ways of tried to get this working and I always get the same error. I can’t seem to find any clue as to what’s going on in Lightning’s logging docs. Anyone know what’s going on? Or at least how to properly log something like this.

Found the solution was to use self.logger.experiment.log