from torch import randn
from torchmetrics.regression import CosineSimilarity
metric = CosineSimilarity()
metric.update(randn(10,), randn(10,))
fig_, ax_ = metric.plot()