Fix data_iter in prepare_dataset from speechcommands example (ml-expl…

…ore#1113)
mokeddembillel · Dec 16, 2024 · a73de93 · a73de93
1 parent a0e7965
commit a73de93
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/speechcommands/main.py b/speechcommands/main.py
@@ -76,6 +76,7 @@ def step(x, y):
     samples_per_sec = []
 
     model.train(True)
+    train_iter.reset()
     for batch_counter, batch in enumerate(train_iter):
         x = mx.array(batch["audio"])
         y = mx.array(batch["label"])
@@ -111,6 +112,7 @@ def test_epoch(model, test_iter):
     model.train(False)
     accs = []
     throughput = []
+    test_iter.reset()
     for batch_counter, batch in enumerate(test_iter):
         x = mx.array(batch["audio"])
         y = mx.array(batch["label"])