I'm going to try and implement this architecture: https://arxiv.org/pdf/1410.5401.pdf
I've taken a lot of inspiration from: https://github.com/snipsco/ntm-lasagne
This will be used for other work in the future too so hopefully it's extensible and able to be used in other projects.
I've just coded something neat that I'll be working on from here forward using this code - my new pytorch-EMM project. It's basically a project for just the external memory part of the NTM, which allows us to unlock the updates (if we wanted to) or do other cool stuff!