one feedback talking to a few researchers and showing the design is that the optim API right now doesn't allow one to optimize a non parameter variable.
For example, when optimizing like neural-art maybe, but also a bunch of meta-optimization research (like learning to learning to learn).
We should consider making it optional to give model in the constructor, and allow .step to take in params to optimize