Repository: Lightning-AI/torchmetrics

Issue:
IndexError when computing multiclass accuracy with topk and ignore_index
## 🐛 Bug

When trying to compute multiclass accuracy with topk set to a value greater than 1, as well as an ignore index provided, torchmetrics throws an IndexError, failing to compute the accuracy. 

### To Reproduce

```python
from torchmetrics.functional import accuracy
import torch

n_batch = 12
n_class = 1024
n_seq = 200

# setup dummy predictions and ground truth
preds = torch.rand((n_batch, n_class, n_seq))
true = torch.randint(0, n_class-1, (n_batch, n_seq))

# mask a portion of our targers
true[:, :, 0:10] = -100

accuracy(
    preds, true, 
    task='multiclass', 
    num_classes=n_class, 
    ignore_index=-100, 
    top_k=10
  )
```

```
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
[<ipython-input-3-c9c955eb3194>](https://localhost:8080/#) in <module>
      8 true = torch.randint(0, n_class-1, (n_batch, n_seq))
      9 
---> 10 accuracy(preds, true, task='multiclass', num_classes=n_class, ignore_index=-100, top_k=10)

2 frames
[/usr/local/lib/python3.8/dist-packages/torchmetrics/functional/classification/stat_scores.py](https://localhost:8080/#) in _multiclass_stat_scores_update(preds, target, num_classes, top_k, average, multidim_average, ignore_index)
    367             target = target.clone()
    368             idx = target == ignore_index
--> 369             preds[idx] = num_classes
    370             target[idx] = num_classes
    371 

IndexError: The shape of the mask [12, 200] at index 1 does not match the shape of the indexed tensor [12, 1024, 200] at index 1
```


### Expected behavior

The function should compute topk accuracy, and ignore the indices specified. 

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source): 0.11.1, installed via pip
- Python & PyTorch Version (e.g., 1.0):  torch 1.13.1+cu116, Python 3.8.10
- Any other relevant information such as OS (e.g., Linux): Colab notebook, but also repros on Linux





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

Resolving patch:
diff --git a/src/torchmetrics/functional/classification/stat_scores.py b/src/torchmetrics/functional/classification/stat_scores.py
--- a/src/torchmetrics/functional/classification/stat_scores.py
+++ b/src/torchmetrics/functional/classification/stat_scores.py
@@ -358,8 +358,9 @@ def _multiclass_stat_scores_update(
             preds = preds.clone()
             target = target.clone()
             idx = target == ignore_index
-            preds[idx] = num_classes
             target[idx] = num_classes
+            idx = idx.unsqueeze(1).repeat(1, num_classes, 1) if preds.ndim > target.ndim else idx
+            preds[idx] = num_classes
 
         if top_k > 1:
             preds_oh = torch.movedim(select_topk(preds, topk=top_k, dim=1), 1, -1)
@@ -374,7 +375,7 @@ def _multiclass_stat_scores_update(
             if 0 <= ignore_index <= num_classes - 1:
                 target_oh[target == ignore_index, :] = -1
             else:
-                preds_oh = preds_oh[..., :-1]
+                preds_oh = preds_oh[..., :-1] if top_k == 1 else preds_oh
                 target_oh = target_oh[..., :-1]
                 target_oh[target == num_classes, :] = -1
         sum_dim = [0, 1] if multidim_average == "global" else [1]


Test patch:
diff --git a/tests/unittests/classification/test_stat_scores.py b/tests/unittests/classification/test_stat_scores.py
--- a/tests/unittests/classification/test_stat_scores.py
+++ b/tests/unittests/classification/test_stat_scores.py
@@ -343,6 +343,21 @@ def test_top_k_multiclass(k, preds, target, average, expected):
     )
 
 
+def test_top_k_ignore_index_multiclass():
+    """Test that top_k argument works together with ignore_index."""
+    preds_without = torch.randn(10, 3).softmax(dim=-1)
+    target_without = torch.randint(3, (10,))
+    preds_with = torch.cat([preds_without, torch.randn(10, 3).softmax(dim=-1)], 0)
+    target_with = torch.cat([target_without, -100 * torch.ones(10)], 0).long()
+
+    res_without = multiclass_stat_scores(preds_without, target_without, num_classes=3, average="micro", top_k=2)
+    res_with = multiclass_stat_scores(
+        preds_with, target_with, num_classes=3, average="micro", top_k=2, ignore_index=-100
+    )
+
+    assert torch.allclose(res_without, res_with)
+
+
 def test_multiclass_overflow():
     """Test that multiclass computations does not overflow even on byte input."""
     preds = torch.randint(20, (100,)).byte()