From 0ca162cfb2a85164cbec18304913c4220520786c Mon Sep 17 00:00:00 2001 From: sakares saengkaew Date: Tue, 3 Dec 2024 14:56:07 +0700 Subject: [PATCH] Fix data_iter in prepare_dataset from speechcommands example (#1113) --- speechcommands/main.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/speechcommands/main.py b/speechcommands/main.py index 0d8da9fd..ed328f4c 100644 --- a/speechcommands/main.py +++ b/speechcommands/main.py @@ -76,6 +76,7 @@ def train_epoch(model, train_iter, optimizer, epoch): 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"])