In order to benefit from `.to(device)` https://discuss.pytorch.org/t/nn-moduledict-vs-ordereddict-vs-dict/145418