Thanks, I'm glad you liked it! Mode collapse was actually the one thing I never encountered during my exploration (which was the reason we looked into using GTNs as a mode-collapse solution for GANs). That said, I found meta-learning to be surprisingly hard to implement efficiently and ran into more bugs in both PyTorch and TensorFlow than I can count.

Changing the learner architecture is not that important actually so that's probably not your problem.