Validation is only executed once in 02_mnist_train_loop.py example

[mstoelzle]

If I read the script examples/flax/02_mnist_train_loop.py correctly, the test loop should be executed in total 10 times. However, after inserting a debug statement into the test_step method, I identified that the test loop is in reality only executed once (at the very beginning of the training loop).

I suspect that this buggy behavior is connected to some kind of faulty loading/cut-off of the validation dataset iterator:

The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset  will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead.

I would appreciate your help in fixing this wrong behaviour.

[cgarciae]

Hey @mstoelzle, thanks for reporting this! I have confirmed this. The problem is that the second time the test_loop runs all of its tasks are empty. Fixing...

[cgarciae]

Fixed in 0.1.8