import torch
from torchmetrics.nominal import TheilsU
metric = TheilsU(num_classes=10)
metric.update(torch.randint(10, (10,)), torch.randint(10, (10,)))
fig_, ax_ = metric.plot()