Modular RNN in PyTorch