Just letting everyone know that following Thomas’s suggestion in another thread, I whipped up a JIT-based implementation a while back that works much faster than using nn.Linear.
It’s available here: https://github.com/Maghoumi/JitGRU 31