In case it might be useful for someone, by modifying the previous code in this way works in my case (the previous code worked on Colab but not locally in my case)
import matplotlib.pyplot as plt
plt.switch_backend('agg')

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)

writer.add_figure('my_fig', fig)
writer.close()


Now the confusion matrix is correctly written on Tensorboard.