Repository: deepset-ai/haystack

Issue:
`FARMReader.train()` raises a torch error
**Describe the bug**
`FARMReader.train` raises a torch error in the [tutorial](https://colab.research.google.com/github/deepset-ai/haystack-tutorials/blob/main/tutorials/02_Finetune_a_model_on_your_data.ipynb) run on colab.

```python
from haystack.nodes import FARMReader
from haystack.utils import fetch_archive_from_http

reader = FARMReader(model_name_or_path="distilbert-base-uncased-distilled-squad", use_gpu=True)
data_dir = "data/squad20"
reader.train(data_dir=data_dir, train_filename="dev-v2.0.json", use_gpu=True, n_epochs=1, save_dir="my_model")
```

**Error message**
```python
RuntimeError: stack expects each tensor to be equal size, but got [6, 2] at entry 0 and [5, 2] at entry 1` error
```

**Expected behavior**
No error

**Additional context**
Discussed in https://github.com/deepset-ai/haystack/discussions/5081
Works fine with Haystack `v1.16.1`. 

**To Reproduce**
* Run tutorial https://haystack.deepset.ai/tutorials/02_finetune_a_model_on_your_data
or
* The code snippet above

**FAQ Check**
- [x] Have you had a look at [our new FAQ page](https://docs.haystack.deepset.ai/docs/faq)?

**System:**
 - OS: Google Colab
 - GPU/CPU: GPU
 - Haystack version (commit or version number): 1.17.1
 - DocumentStore: - 
 - Reader: FARMReader
 - Retriever: -


Discussion:
Once this issue is solved, the version in tutorial 2 should be updated. This is tracked with this issue -> https://github.com/deepset-ai/haystack-tutorials/issues/202

Resolving patch:
diff --git a/haystack/modeling/data_handler/dataloader.py b/haystack/modeling/data_handler/dataloader.py
--- a/haystack/modeling/data_handler/dataloader.py
+++ b/haystack/modeling/data_handler/dataloader.py
@@ -3,6 +3,7 @@
 from math import ceil
 
 import torch
+import torch.nn.functional as F
 from torch.utils.data import DataLoader, Dataset, Sampler
 
 from haystack.errors import ModelingError
@@ -40,18 +41,26 @@ def collate_fn(batch):
             else:
                 _tensor_names = tensor_names
 
-            if type(batch[0]) == list:
+            if isinstance(batch[0], list):
                 batch = batch[0]
 
             if len(batch[0]) != len(_tensor_names):
                 raise ModelingError(
                     f"Dataset contains {len(batch[0])} tensors while there are {len(_tensor_names)} tensor names supplied: {_tensor_names}"
                 )
-            lists_temp = [[] for _ in range(len(_tensor_names))]
-            ret = dict(zip(_tensor_names, lists_temp))
 
+            max_num_labels = self._compute_max_number_of_labels(batch=batch, tensor_names=_tensor_names)
+
+            ret = {name: [] for name in tensor_names}
             for example in batch:
                 for name, tensor in zip(_tensor_names, example):
+                    # each example may have a different number of answers/labels,
+                    # so we need to pad the corresponding tensors to the max number of labels
+                    if name == "labels" and tensor.ndim > 0:
+                        num_labels = tensor.size(0)
+                        if num_labels < max_num_labels:
+                            padding = (0, 0, 0, max_num_labels - num_labels)
+                            tensor = F.pad(tensor, padding, value=-1)
                     ret[name].append(tensor)
 
             for key in ret:
@@ -75,3 +84,15 @@ def __len__(self):
             return num_batches
         else:
             return super().__len__()
+
+    def _compute_max_number_of_labels(self, batch, tensor_names) -> int:
+        """
+        Compute the maximum number of labels in a batch.
+        Each example may have a different number of labels, depending on the number of answers.
+        """
+        max_num_labels = 0
+        for example in batch:
+            for name, tensor in zip(tensor_names, example):
+                if name == "labels" and tensor.ndim > 0:
+                    max_num_labels = max(max_num_labels, tensor.size(0))
+        return max_num_labels


Test patch:
diff --git a/test/modeling/test_dataloader.py b/test/modeling/test_dataloader.py
new file mode 100644
--- /dev/null
+++ b/test/modeling/test_dataloader.py
@@ -0,0 +1,47 @@
+import pytest
+
+import torch
+
+from haystack.modeling.data_handler.dataloader import NamedDataLoader
+
+
+@pytest.fixture
+def named_dataloader():
+    tensor_names = ["input_ids", "labels"]
+    return NamedDataLoader(None, 1, tensor_names=tensor_names)
+
+
+@pytest.fixture
+def batch():
+    # batch containing tensors of different lengths
+    return [
+        (torch.tensor([1, 2, 3]), torch.tensor([[0, 0], [-1, -1], [-1, -1], [-1, -1], [-1, -1], [-1, -1]])),
+        (torch.tensor([4, 5, 6]), torch.tensor([[0, 0], [-1, -1], [-1, -1]])),
+        (torch.tensor([7, 8, 9]), torch.tensor([[0, 0], [-1, -1]])),
+    ]
+
+
+@pytest.mark.unit
+def test_compute_max_number_of_labels(named_dataloader, batch):
+    tensor_names = ["input_ids", "labels"]
+    max_num_labels = named_dataloader._compute_max_number_of_labels(batch, tensor_names)
+    assert max_num_labels == 6
+
+
+@pytest.mark.unit
+def test_collate_fn(named_dataloader, batch):
+    collated_batch = named_dataloader.collate_fn(batch)
+
+    expected_collated_batch = {
+        "input_ids": torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
+        "labels": torch.tensor(
+            [
+                [[0, 0], [-1, -1], [-1, -1], [-1, -1], [-1, -1], [-1, -1]],
+                [[0, 0], [-1, -1], [-1, -1], [-1, -1], [-1, -1], [-1, -1]],
+                [[0, 0], [-1, -1], [-1, -1], [-1, -1], [-1, -1], [-1, -1]],
+            ]
+        ),
+    }
+
+    for key in collated_batch:
+        assert torch.equal(collated_batch[key], expected_collated_batch[key])