num_samples向下去整,防止数据集的溢出 (PaddlePaddle#8691)

DesmonDay · Sep 5, 2024 · a55589f · a55589f
1 parent 6e24524
commit a55589f
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/paddlenlp/utils/batch_sampler.py b/paddlenlp/utils/batch_sampler.py
@@ -14,8 +14,6 @@
 
 from __future__ import division, print_function
 
-import math
-
 import paddle
 
 __all__ = ["DistributedBatchSampler"]
@@ -110,7 +108,7 @@ def __init__(
             # In pre-training mode when using distributed dataloader, the input dataset can be None. We should handle this situation.
             self.num_samples = 0
         else:
-            self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.nranks))
+            self.num_samples = int(len(self.dataset) * 1.0 / self.nranks)
         self.total_size = self.num_samples * self.nranks
 
     def get_start_end_idx(self):
@@ -125,7 +123,7 @@ def __iter__(self):
             self.consumed_samples,
             self.nranks,
         )
-        self.remain_num_samples = int(math.ceil((len(self.dataset) - self.consumed_samples) * 1.0 / self.nranks))
+        self.remain_num_samples = int((len(self.dataset) - self.consumed_samples) * 1.0 / self.nranks)
         self.remain_total_size = self.remain_num_samples * self.nranks
         self.batch_size_times_rank_size = self.batch_size * self.nranks