diff --git a/Compiler/ml.py b/Compiler/ml.py
index 3cd738468..417f440a0 100644
--- a/Compiler/ml.py
+++ b/Compiler/ml.py
@@ -689,7 +689,7 @@ def _(j):
                         indices = indices_by_label[label]
                         @for_range_multithread(self.n_threads, 1, n)
                         def _(i):
-                            idx = indices[i + j * n_per_epoch]
+                            idx = indices[i + j * n]
                             self.layers[0].X[i + label * n] = X[idx]
                     self.forward(None)
                     self.backward()