About training code

@lukas-schott , thank you for your implementation, can you also show the training code in order to me to evaluate it on other datasets?

Thanks

Review the repository's existing implementation to identify which training code is missing and where the current model workflow begins. Determine what would let someone evaluate the implementation on other datasets; done means the training procedure is available and usable beyond the provided MNIST setup.

Written by the indexing model from the issue text.