MulticlassAccuracy Calculation Results in Dimension Mismatch Error

[okezieowen]

When I try to compute the multiclass accuracy for a preds/target input with a top_k value greater than 1, I get an error about mismatched dimensions.
According to the documentation, preds can be a float tensor of shape (N, C, ...) and target an integer tensor of shape (N, ...).

To Reproduce

Steps to reproduce the behavior...

1. Create a preds floating tensor of shape (batch_size, num_classes, sequence_length)
2. Create a target tensor of shape (batch_size, sequence_length)
3. Instantiate MulticlassAccuracy with a top_k > 1 and multidim_average set to global.
4. Call the forward method, using preds and target as the inputs.

I get the error message:
RuntimeError: The size of tensor a (4) must match the size of tensor b (7) at non-singleton dimension 0

Code sample

# Ideally attach a minimal code sample to reproduce the decried issue.
# Minimal means having the shortest code but still preserving the bug.
!pip install torch torchmetrics --quiet --upgrade

import torch
import torchmetrics

batch_size = 1
num_classes = 4
sequence_length = 7
top_k = 3

preds = torch.randn(batch_size, num_classes, sequence_length)
targets = torch.randint(num_classes, (batch_size, sequence_length))

accuracy = MulticlassAccuracy(num_classes = num_classes, top_k = top_k, multidim_average='global')
acc = accuracy(preds, targets)

Expected Behavior

I expect the output to be a scalar tensor as mentioned in the documentation.

Environment

• TorchMetrics version 1.6.1
• Python version: 3.11.11
• PyTorch Version: 2.6.0+cu124

Additional context