I am using the code inside the on_train_epoch_end inside a metric. So the important part is:
` fig = plt.figure()

my plot

logger.experiment.add_figure("fig", fig)
plt.close() `