Hi @wwc, I was answering Sagar’s question down below in the comments and I sent him a sample code that trains on MNIST dataset using the lib we just built. I thought you might be interested, so here is the link to the answer :

https://medium.com/@omaraflak/hey-sagar-i-updated-the-code-of-the-flatten-layer-i-made-a-small-mistake-a7a5fe9c5e8f