Update dataloader.py

Tencent · Nov 5, 2023 · 4597691 · 4597691
1 parent 3ab1161
commit 4597691
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tencentpretrain/utils/dataloader.py b/tencentpretrain/utils/dataloader.py
@@ -28,6 +28,11 @@ def __init__(self, args, dataset_path, batch_size, global_rank, world_size, loca
         self.span_masking = args.span_masking
         self.span_geo_prob = args.span_geo_prob
         self.span_max_length = args.span_max_length
+        self.skip_data_num = args.skip_data_num
+
+        if self.skip_data_num > 0:
+            for _ in range(self.skip_data_num):
+                instance = pickle.load(self.dataset_reader)
 
     def _fill_buf(self):
         try: