For comparison, here's a solution in torch with a torchscript derocator: https://gist.github.com/bethebunny/1319b96e412d6a68488b7cb35...
IMO it's a huge advantage to be able to write torch model code that looks like how you'd write the same program in pure python, and still have them be serializable.