Repository: Lightning-AI/torchmetrics

Issue:
Incorrect split of columns and rows in a plot for multi label confusion matrix.
## 🐛 Bug

The specific values of `num_labels` in `MultilabelConfusionMatrix` can throw `IndexError` in the `plot` function. 

The `_get_col_row_split` function for `num_labels=11` returns `rows=3, cols=3`, which is not enough for 11 axes.

### To Reproduce

```Python
import torch
from torchmetrics.classification import MultilabelConfusionMatrix

num_labels = 11  # 17 21
example_preds = torch.rand((100, num_labels))
example_targets = torch.randint(0, 2, (100, num_labels))

conf_matrix = MultilabelConfusionMatrix(num_labels)
conf_matrix(example_preds, example_targets)
conf_matrix.plot()
```

<details>
  <summary>Traceback</summary>

```
Traceback (most recent call last):
  File "D:\test\test.py", line 10, in <module>
    conf_matrix.plot()
  File "D:\test\.venv\lib\site-packages\torchmetrics\classification\confusion_matrix.py", line 466, in plot
    fig, ax = plot_confusion_matrix(val, ax=ax, add_text=add_text, labels=labels)
  File "C:\Users\Anton\AppData\Local\Programs\Python\Python310\lib\contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "D:\test\.venv\lib\site-packages\torchmetrics\utilities\plot.py", line 248, in plot_confusion_matrix
    ax = axs[i] if rows != 1 and cols != 1 else axs
IndexError: index 9 is out of bounds for axis 0 with size 9
```
<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->

</details>

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source): 1.1.1, pip
- Python & PyTorch Version (e.g., 1.0): 3.10.9, 2.0.1
- Any other relevant information such as OS (e.g., Linux): Windows




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

Resolving patch:
diff --git a/src/torchmetrics/utilities/plot.py b/src/torchmetrics/utilities/plot.py
--- a/src/torchmetrics/utilities/plot.py
+++ b/src/torchmetrics/utilities/plot.py
@@ -172,9 +172,9 @@ def plot_single_or_multi_val(
 def _get_col_row_split(n: int) -> Tuple[int, int]:
     """Split `n` figures into `rows` x `cols` figures."""
     nsq = sqrt(n)
-    if nsq * nsq == n:
+    if int(nsq) == nsq:  # square number
         return int(nsq), int(nsq)
-    if floor(nsq) * ceil(nsq) > n:
+    if floor(nsq) * ceil(nsq) >= n:
         return floor(nsq), ceil(nsq)
     return ceil(nsq), ceil(nsq)
 


Test patch:
diff --git a/tests/unittests/utilities/test_plot.py b/tests/unittests/utilities/test_plot.py
--- a/tests/unittests/utilities/test_plot.py
+++ b/tests/unittests/utilities/test_plot.py
@@ -168,6 +168,7 @@
     _TORCH_GREATER_EQUAL_1_10,
     _TORCHAUDIO_GREATER_EQUAL_0_10,
 )
+from torchmetrics.utilities.plot import _get_col_row_split
 from torchmetrics.wrappers import (
     BootStrapper,
     ClasswiseWrapper,
@@ -794,6 +795,17 @@ def test_plot_methods_retrieval(metric_class, preds, target, indexes, num_vals):
     plt.close(fig)
 
 
+@pytest.mark.parametrize(
+    ("n", "expected_row", "expected_col"),
+    [(1, 1, 1), (2, 1, 2), (3, 2, 2), (4, 2, 2), (5, 2, 3), (6, 2, 3), (7, 3, 3), (8, 3, 3), (9, 3, 3), (10, 3, 4)],
+)
+def test_row_col_splitter(n, expected_row, expected_col):
+    """Test the row col splitter function works as expected."""
+    row, col = _get_col_row_split(n)
+    assert row == expected_row
+    assert col == expected_col
+
+
 @pytest.mark.parametrize(
     ("metric_class", "preds", "target", "labels"),
     [