Implementing a metric is getting complicated
In PL 0.9.0 I had the Metric implementation bellow:
class MRRMetric(TensorMetric):

    def similarities(self, x1, x2):
        """
        Calculates the cosine similarity matrix for every pair (i, j),
        where i is an embedding from x1 and j is another embedding from x2.

        :param x1: a tensors with shape [batch_size, hidden_size].
        :param x2: a tensors with shape [batch_size, hidden_size].
        :return: the cosine similarity matrix with shape [batch_size, batch_size].
        """
        x1 = x1 / torch.norm(x1, dim=1, p=2, keepdim=True, dtype=torch.float32)
        x2 = x2 / torch.norm(x2, dim=1, p=2, keepdim=True, dtype=torch.float32)
        return torch.matmul(x1, x2.t())

    def forward(self, r1, r2):
        distances = 1 - self.similarities(r1, r2)
        correct_elements = torch.unsqueeze(torch.diag(distances), dim=-1)
        batch_ranks = torch.sum(distances < correct_elements, dim=-1) + 1.0

        return torch.mean(1.0 / batch_ranks)
In PL 1.0.8 the same metric is implemented with the same logic as:
class MRRMetric(Metric):
    def __init__(self):
        super().__init__()
        self.add_state("mrrs", default=[])

    def similarities(self, x1, x2):
        """
        Calculates the cosine similarity matrix for every pair (i, j),
        where i is an embedding from x1 and j is another embedding from x2.

        :param x1: a tensors with shape [batch_size, hidden_size].
        :param x2: a tensors with shape [batch_size, hidden_size].
        :return: the cosine similarity matrix with shape [batch_size, batch_size].
        """
        x1 = x1 / torch.norm(x1, dim=1, p=2, keepdim=True, dtype=torch.float32)
        x2 = x2 / torch.norm(x2, dim=1, p=2, keepdim=True, dtype=torch.float32)
        return torch.matmul(x1, x2.t())

    def update(self, r1, r2):
        distances = 1 - self.similarities(r1, r2)
        correct_elements = torch.unsqueeze(torch.diag(distances), dim=-1)
        batch_ranks = torch.sum(distances < correct_elements, dim=-1) + 1.0
        mrr = torch.mean(1.0 / batch_ranks)
        self.mrrs.append(mrr )
    
    def compute(self):
        return torch.mean(torch.tensor(self.mrrs))
However I'm getting differents outputs when calling this metrics:

PL 0.9.0

s=torch.tensor([[0.0023, 0.8380, 0.2920, 0.2263, 0.3362, 0.7074, 0.5290, 0.5974, 0.8176,
         0.5816],
        [0.6414, 0.9122, 0.0162, 0.3203, 0.5081, 0.9035, 0.3201, 0.2870, 0.0693,
         0.3736],
        [0.4570, 0.2787, 0.0140, 0.5623, 0.7226, 0.3223, 0.3566, 0.6783, 0.8384,
         0.1610],
        [0.1334, 0.7529, 0.5574, 0.1949, 0.9526, 0.9542, 0.2212, 0.1759, 0.1354,
         0.8739],
        [0.3224, 0.2001, 0.2776, 0.4203, 0.0720, 0.2702, 0.0114, 0.1538, 0.0090,
         0.4954]])
r=torch.tensor([[1.7271e-01, 6.2802e-01, 1.9287e-01, 3.7181e-01, 9.8512e-01, 3.3500e-01,
         7.8245e-01, 8.2903e-01, 3.8158e-01, 9.0951e-01],
        [1.6962e-01, 1.8394e-01, 8.1776e-01, 1.5577e-01, 2.6253e-01, 8.9279e-02,
         4.3836e-01, 9.3631e-01, 9.9012e-01, 7.0277e-01],
        [6.0912e-01, 3.9780e-04, 2.9773e-02, 3.0968e-01, 9.5504e-01, 5.9606e-01,
         7.1733e-01, 5.7158e-01, 7.0587e-01, 2.0964e-01],
        [5.1019e-01, 3.9333e-01, 8.7822e-02, 8.1885e-01, 3.9182e-01, 6.6586e-02,
         1.6968e-01, 1.6206e-01, 6.1676e-01, 1.8388e-01],
        [8.3266e-01, 5.5918e-01, 6.3568e-01, 5.4674e-01, 3.4784e-01, 6.0922e-01,
         1.3884e-01, 2.3742e-01, 9.0644e-01, 2.3266e-01]])
mrr = MRRMetric()
mrr(s, r)
#tensor(0.5300)



PL 1.0.8



```python
s=torch.tensor([[0.0023, 0.8380, 0.2920, 0.2263, 0.3362, 0.7074, 0.5290, 0.5974, 0.8176,
         0.5816],
        [0.6414, 0.9122, 0.0162, 0.3203, 0.5081, 0.9035, 0.3201, 0.2870, 0.0693,
         0.3736],
        [0.4570, 0.2787, 0.0140, 0.5623, 0.7226, 0.3223, 0.3566, 0.6783, 0.8384,
         0.1610],
        [0.1334, 0.7529, 0.5574, 0.1949, 0.9526, 0.9542, 0.2212, 0.1759, 0.1354,
         0.8739],
        [0.3224, 0.2001, 0.2776, 0.4203, 0.0720, 0.2702, 0.0114, 0.1538, 0.0090,
         0.4954]])
r=torch.tensor([[1.7271e-01, 6.2802e-01, 1.9287e-01, 3.7181e-01, 9.8512e-01, 3.3500e-01,
         7.8245e-01, 8.2903e-01, 3.8158e-01, 9.0951e-01],
        [1.6962e-01, 1.8394e-01, 8.1776e-01, 1.5577e-01, 2.6253e-01, 8.9279e-02,
         4.3836e-01, 9.3631e-01, 9.9012e-01, 7.0277e-01],
        [6.0912e-01, 3.9780e-04, 2.9773e-02, 3.0968e-01, 9.5504e-01, 5.9606e-01,
         7.1733e-01, 5.7158e-01, 7.0587e-01, 2.0964e-01],
        [5.1019e-01, 3.9333e-01, 8.7822e-02, 8.1885e-01, 3.9182e-01, 6.6586e-02,
         1.6968e-01, 1.6206e-01, 6.1676e-01, 1.8388e-01],
        [8.3266e-01, 5.5918e-01, 6.3568e-01, 5.4674e-01, 3.4784e-01, 6.0922e-01,
         1.3884e-01, 2.3742e-01, 9.0644e-01, 2.3266e-01]])
mrr = MRRMetric()
mrr(r, s)
#tensor(0.6800)