From c01e65bb48d8f0b924f31ff627b257e73fd457d1 Mon Sep 17 00:00:00 2001
From: Xueqing Liu <liususan091219@users.noreply.github.com>
Date: Sun, 9 Oct 2022 00:59:31 -0400
Subject: [PATCH] updating the data collator for seq-regression to handle the
 dim mismatch problem (#751)

---
 flaml/nlp/huggingface/data_collator.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/flaml/nlp/huggingface/data_collator.py b/flaml/nlp/huggingface/data_collator.py
index 2d10f1520..b3911c291 100644
--- a/flaml/nlp/huggingface/data_collator.py
+++ b/flaml/nlp/huggingface/data_collator.py
@@ -11,6 +11,7 @@ from flaml.data import (
     MULTICHOICECLASSIFICATION,
     SUMMARIZATION,
     SEQCLASSIFICATION,
+    SEQREGRESSION
 )
 
 
@@ -51,5 +52,6 @@ task_to_datacollator_class = OrderedDict(
         (MULTICHOICECLASSIFICATION, DataCollatorForMultipleChoiceClassification),
         (SUMMARIZATION, DataCollatorForSeq2Seq),
         (SEQCLASSIFICATION, DataCollatorWithPadding),
+        (SEQREGRESSION, DataCollatorWithPadding),
     ]
 )