Calculations nDCG using GPU are 2x slower than CPU #2287

import timeit

import numpy as np
import torch
from sklearn.metrics import ndcg_score
from torchmetrics.functional.retrieval import retrieval_normalized_dcg

# p and t are examples given by both sklearn and torchmetrics
p = [.1, .2, .3, 4, 70] * 100
t = [10, 0, 0, 1, 5] * 100

number = int(1e4)

# 1. BENCHMARK: numpy array
preds = np.asarray([p])
target = np.asarray([t])

def a():
    return ndcg_score(target, preds)

print(f'numpy array: {timeit.timeit("a()", setup="from __main__ import a", number=number):.4f}')

# 2. cpu tensor
preds_cpu = torch.tensor(p)
target_cpu = torch.tensor(t)

assert preds_cpu.device == torch.device("cpu")

def b():
    retrieval_normalized_dcg(preds_cpu, target_cpu)

print(f'CPU tensor: {timeit.timeit(f"b()", setup="from __main__ import b", number=number):.4f}')

# 3. gpu tensor
preds_gpu = torch.tensor(p, device="cuda")
target_gpu = torch.tensor(t, device="cuda")

assert preds_gpu.device == torch.device("cuda:0")

def c():
    retrieval_normalized_dcg(preds_gpu, target_gpu)

print(f'GPU tensor: {timeit.timeit("c()", setup="from __main__ import c", number=number):.4f}')