Comment by danielhanchen

Hey there! Yes that's exactly what I thought! I was writing up a blog post at https://colab.research.google.com/drive/1AOuhMVILE06mD-Go7-R... which step by step shows every change I did plus gave timings / mem savings.

I'll post it once its completed if you're interested!