fix(training): reinitialize dataset each epoch to avoid iterator exhaustion

[Shocker444]

Description

This is a fix (although I assume temporary) to the training scheme which solves the issue of dataset exhaustion after the first epoch.

Currently the Kaun.iter method exhausts the dataset, yielding nothing the second epoch which returns no loss and ultimately no accuracy('Invalid accuracy state)

Fix

• This fix re initializes the dataset and performs all necessary manipulations(shuffle, batch, etc) at the start of each epoch, helps to prevent exhaustion
• Temporary fix as initializing dataset at each epoch may create some overhead.

Results

• mnist_ml.exe trains for complete 10 epochs without "invalid accuracy state" error
• All tests pass dune runtest kaun