# pip install --upgrade torch
# pip install tensorboard

import matplotlib.pyplot as plt

import tensorboard
import torch

from sklearn.metrics import confusion_matrix
from torch.utils.tensorboard import SummaryWriter

print(torch.__version__) # 1.5.1
print(tensorboard.__version__)  # 2.2.2

writer = SummaryWriter('runs/experiment')

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))

fig = plt.figure()
plt.imshow(conf_mat)
plt.show()

writer.add_figure('my_fig', fig)

And here’s the tensorboard snapshot Screen Shot 2020-06-24 at 11.45.25 AM1275×508 33.4 KB