Fix tests

PetrochukM · Apr 8, 2018 · 1424e80 · 1424e80
1 parent dc6e144
commit 1424e80
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/torchnlp/samplers/bucket_batch_sampler.py b/torchnlp/samplers/bucket_batch_sampler.py
@@ -1,5 +1,6 @@
 import heapq
 import pickle
+import math
 
 from torch.utils.data.sampler import BatchSampler
 from torch.utils.data.sampler import RandomSampler
@@ -114,4 +115,4 @@ def __len__(self):
         if self.drop_last:
             return len(self.data) // self.batch_size
         else:
-            return len(self.data) // self.batch_size + 1
+            return math.ceil(len(self.data) / self.batch_size)
diff --git a/torchnlp/samplers/noisy_sorted_batch_sampler.py b/torchnlp/samplers/noisy_sorted_batch_sampler.py
@@ -74,4 +74,4 @@ def __iter__(self):
             random.shuffle(batches)
         if self.last_batch_first:
             batches.insert(0, last_batch)
-        return iter(batches)
+        return iter(batches)