From a524d01ac3a584debc8edc7bdb7f6c29e80fe34c Mon Sep 17 00:00:00 2001
From: Colin <colin>
Date: Wed, 17 Apr 2024 23:59:09 +0800
Subject: [PATCH] Add multithread dataloder support.

---
 wit/meaning_dataset.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/wit/meaning_dataset.py b/wit/meaning_dataset.py
index b3c89fc..b60be4c 100644
--- a/wit/meaning_dataset.py
+++ b/wit/meaning_dataset.py
@@ -447,6 +447,14 @@ class BatchGroupMeaningDataloader(Dataset):
             s += "--------------------------------------------------------\n"
         return s
 
+    def detection_collate(batch):
+        return batch[0]
+
+    def dataloader(self, num_workers=1):
+        return DataLoader(
+            self, batch_size=1, num_workers=num_workers, collate_fn=BatchGroupMeaningDataloader.detection_collate
+        )
+
 
 if __name__ == "__main__":