Too slow on cpu.

[hihunjin]

🐛 Bug

Using torchmetrics with cpu is too slow, slower than with gpu. I don't think there are a lot of calculations using gpu, multiplies, divides. But I'm not sure where this bottleneck took.

To Reproduce

Steps to reproduce the behavior:

1. Go to 'colab'

import torch
import torchmetrics

accuracy_train = torchmetrics.Accuracy(num_classes=5, average='none')
x = torch.rand((1,5,40,512,512))
y = torch.randint(0,5,(1,40,512,512))

## to gpu
x_ = x.cuda()
y_ = y.cuda()
accuracy_train.cuda()

%%timeit
accuracy_train(x_,y_)

accuracy_train.cpu()

%%timeit
accuracy_train(x,y)

2. See the following outcomes.

Expected behavior

They should not be that much different.

Environment

• google colab

[github-actions]

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

[SkafteNicki]

After running the python cprofiler and vizualizing with snakeviz I get this:

it is the topk function from pytorch that is the sinner here, that seems to have much worse performance on cpu than gpu (the program spends around 92% on that single call).
We could probably work around this to only use the torch.topk function whenever the topk parameter != 1 (just use argmax which I think have better performance). @Borda any opinion here?