add drop_last

konas122 · Mar 7, 2024 · cb88bad · cb88bad
1 parent 4c053bb
commit cb88bad
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/dazero/dataloaders.py b/dazero/dataloaders.py
@@ -10,14 +10,19 @@
 
 
 class DataLoader:
-    def __init__(self, dataset, batch_size, shuffle=True, gpu=False):
+    def __init__(self, dataset, batch_size, shuffle=True, gpu=False, drop_last=False):
         if not isinstance(dataset, Dataset):
             raise TypeError("It must be of type dazero.Dataset")
         self.dataset = dataset
+        self.drop_last = drop_last
         self.batch_size = batch_size
         self.shuffle = shuffle
         self.data_size = len(dataset)
-        self.max_iter = math.ceil(self.data_size / batch_size)
+
+        if drop_last:
+            self.max_iter = self.data_size // batch_size
+        else:
+            self.max_iter = math.ceil(self.data_size / batch_size)
         self.gpu = gpu
 
         self._reset()