There's batch size missing on mnist_continued.py I am not sure, but adding ``` batch_size: 128, ``` in algorithm seems to fix it.