Fix recall metric calculation

[sungho-ham]

Goals ⚽

• Fix recall metric error when batch size is 1

Implementation Details 🚧

Transformers4Rec/transformers4rec/torch/ranking_metric.py

Line 134 in 23d5e3b

rel_indices = (num_relevant != 0).nonzero().squeeze()

The squeeze() function removes all dimensions when batch size is 1.
I've added dim=1 argument to keep dimension along batch.

Testing Details 🔍

Test code for error reproduction

import torch
from transformers4rec.torch import ranking_metric

ks = [1, 2]
scores = torch.tensor( [ [1,2,3], ])
labels = torch.tensor( [ [0,1,0], ])
ranking_metric.RecallAt()._metric(ks, scores, labels)

Error message

    135 rel_indices = (num_relevant != 0).nonzero().squeeze()
    136 rel_count = num_relevant[rel_indices]
--> 138 if rel_indices.shape[0] > 0:
    139     for index, k in enumerate(ks):
    140         rel_labels = topk_labels[rel_indices, : int(k)]

IndexError: tuple index out of range

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.