import torch
from torchmetrics.wrappers import MultioutputWrapper
from torchmetrics.regression import R2Score
metric = MultioutputWrapper(R2Score(), 2)
values = [ ]
for _ in range(3):
    values.append(metric(torch.randn(20, 2), torch.randn(20, 2)))
fig_, ax_ = metric.plot(values)
