Lessons 4 - MNIST with DataBlock API

I tried replicating the MNIST example from Lession 4 with the DataBlock API. I have two problems:

  • The learner takes much longer to train (minutes instead of seconds)
  • The learner does not learn (the accuracy stays around 0.5)

Can someone of you look at it and point out what I am doing wrong. I have a hard time finding the error:

https://colab.research.google.com/drive/1P2ZWScC0Ygm-GqYwJOFM4wY47MHVpsm-#scrollTo=GQ7hlEd-OOHP

Just for reference in the original notebook the whole training of all 10 epochs takes 4 seconds on a CPU node. In my notebook it takes 13 seconds per epoch. A total of 130 seconds, also on a CPU node.