diff --git a/mlpp_lib/datasets.py b/mlpp_lib/datasets.py
index c8b7bfe..cacfbe9 100644
--- a/mlpp_lib/datasets.py
+++ b/mlpp_lib/datasets.py
@@ -593,7 +593,7 @@ def __init__(
         self.shuffle = shuffle
         self.block_size = block_size
         self.num_samples = len(self.dataset.x)
-        self.num_batches = self.num_samples // batch_size
+        self.num_batches = self.num_samples // batch_size if batch_size <= self.num_samples else 1
         self._indices = tf.range(self.num_samples)
         self._seed = 0
         self._reset()