A lot of the hacks to make this work would be achieved much more easily using: https://pytorch.org/functorch/stable/generated/functorch.make_functional_with_buffers.html?highlight=make_functional_with_buffers#functorch.make_functional_with_buffers If I ever come back to rewrite this I would try to do that.
A lot of the hacks to make this work would be achieved much more easily using: https://pytorch.org/functorch/stable/generated/functorch.make_functional_with_buffers.html?highlight=make_functional_with_buffers#functorch.make_functional_with_buffers
If I ever come back to rewrite this I would try to do that.