Hi everyone,
let’s suppose I have this simple code that creates a confusion matrix:
import torch
from sklearn.metrics import confusion_matrix

output = torch.randn(1, 2, 4, 4)
pred = torch.argmax(output, 1)
target = torch.empty(1, 4, 4, dtype=torch.long).random_(2)
conf_mat = confusion_matrix(pred.view(-1), target.view(-1))


Is there a way to convert conf_mat to an image and save it to Tensorboard ?
I successfully converted it to an image using mlxtend library, but now I found no way to save it Tensorboard using  writer.add_image 
import matplotlib.pyplot as plt
from mlxtend.plotting import plot_confusion_matrix

fig, ax = plot_confusion_matrix(conf_mat=conf_mat)
plt.show()


 fig  is of type matplotlib.figure.Figure and it should be converted to a torch.tensor or numpy.array since writer.add_image expects one of these as input. I tried some workarounds but none of them worked.
Any suggestions?