apply shuffle to batch indices

nitrain · Apr 8, 2024 · cad8c93 · cad8c93
1 parent 8823faf
commit cad8c93
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/nitrain/loaders/dataset_loader.py b/nitrain/loaders/dataset_loader.py
@@ -66,7 +66,9 @@ def __iter__(self):
         dataset = self.dataset
         n_image_batches = math.ceil(len(dataset) / batch_size)
 
-        # TODO: shuffle image indices before batch loop
+        original_indices = np.arange(len(dataset))
+        if self.shuffle:
+            np.random.shuffle(original_indices)
 
         image_batch_idx = 0
         while image_batch_idx < n_image_batches: