Perason Correlation with cuda tensors raises error

[jlshin]

When computing the Pearson Correlation Coefficient with input cuda tensors (prediction and target), an error is raised because some default values (mean_x, mean_y) are initialized as torch.tensors (on the cpu).

Example:

    >>> from torchmetrics import PearsonCorrCoef
    >>> target = torch.tensor([3, -0.5, 2, 7]).to('cuda')
    >>> preds = torch.tensor([2.5, 0.0, 2, 8]).to('cuda')
    >>> pearson = PearsonCorrCoef()
    >>> pearson(preds, target)

Output:

python3.9/site-packages/torchmetrics/functional/regression/pearson.py:55, in _pearson_corrcoef_update(preds, target, mean_x, mean_y, var_x, var_y, corr_xy, n_prior)
     53 my_new = (n_prior * mean_y + target.mean() * n_obs) / (n_prior + n_obs)
     54 n_prior += n_obs
---> 55 var_x += ((preds - mx_new) * (preds - mean_x)).sum()
     56 var_y += ((target - my_new) * (target - mean_y)).sum()
     57 corr_xy += ((preds - mx_new) * (target - mean_y)).sum()

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Versions:
torchmetrics: 0.8.2
python3.9

[github-actions]

Hi! thanks for your contribution!, great first issue!

[SkafteNicki]

As described here in the documentation each metric object is a nn.Module and therefore should be moved to the same device as input. Your example therefore works if we move the metric to gpu:

m = PearsonCorrCoef().cuda()
t1 = torch.tensor(range(10), device="cuda", dtype=float)
t2 = 2 * torch.tensor(range(10), device="cuda", dtype=float)
m(t1, t2)

Also we will be improving the error message for this case with this PR: #1056.
Closing issue.