# Using datacollator for multi-task training

**URL:** <https://discuss.huggingface.co/t/using-datacollator-for-multi-task-training/13999>\
**Category:** Intermediate\
**Created:** [January 22, 2022, 6:12pm UTC](https://discuss.huggingface.co/t/using-datacollator-for-multi-task-training/13999 "2022-01-22T18:12:20Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![nid989](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/nid989/32/3579_2.png) [@nid989](https://discuss.huggingface.co/u/nid989)\
**Post date:** [January 22, 2022, 6:12pm UTC](https://discuss.huggingface.co/t/using-datacollator-for-multi-task-training/13999/1 "2022-01-22T18:12:20Z")

</div>

Hey everyone,  
Actually, I am trying to accommodate multiple subtasks including STS-b and NER into a multi-task model, however am unable to issue the tokens from the conll dataset into the DataCollator. Can anyone help me with this? The code snippet is shown below.

```auto
class NLPDataCollator(DefaultDataCollator):
    """
    Extending the existing DataCollator to work with NLP dataset batches
    """
    def collate_batch(self, features: List[Union[InputDataClass, Dict]]) -> Dict[str, torch.Tensor]:
        first = features[0]
        if isinstance(first, dict):
          # NLP data sets current works presents features as lists of dictionary
          # (one per example), so we will adapt the collate_batch logic for that
          if "labels" in first and first["labels"] is not None:
              if first["labels"].dtype == torch.int64:
                  labels = torch.tensor([f["labels"] for f in features], dtype=torch.long)
              else:
                  labels = torch.tensor([f["labels"] for f in features], dtype=torch.float)
              batch = {"labels": labels}
          for k, v in first.items():
              if k != "labels" and v is not None and not isinstance(v, str):
                  batch[k] = torch.stack([f[k] for f in features])
          return batch
        else:
          # otherwise, revert to using the default collate_batch
          return DefaultDataCollator().collate_batch(features)

```

---

<div class="post-metadata">

**Author:** ![nid989](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/nid989/32/3579_2.png) [@nid989](https://discuss.huggingface.co/u/nid989)\
**Post date:** [January 22, 2022, 6:14pm UTC](https://discuss.huggingface.co/t/using-datacollator-for-multi-task-training/13999/2 "2022-01-22T18:14:53Z")

</div>

This error is displayed when I am trying to accommodate the NER dataset specifically entities extracted in the fashion show [here](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/token_classification.ipynb#scrollTo=vc0BSBLIIrJQ)

 ![image](https://us1.discourse-cdn.com/hellohellohello/original/2X/9/9441153db28495c476dfd65f017a340adbef8a44.png)

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](https://discuss.huggingface.co/u/sgugger)\
**Post date:** [January 24, 2022, 2:10pm UTC](https://discuss.huggingface.co/t/using-datacollator-for-multi-task-training/13999/3 "2022-01-24T14:10:49Z")

</div>

It looks like you’re trying to use a very old API for data collators, which was deprecated in v3 and removed in v4. `DefaultDataCollator()` has not `collate_batch` method, you call it directly on the features you want to batch, and you should implement the ` __call__ ` method in your custom data collator.