[Enhancement] Let optimizer use built torch opt

[muellerzr]

This PR slightly tweaks the constructor for OptimWrapper, by allowing it to use an already built pytorch optimizer. This is beneficial for when other libraries need to use a built optimizer directly to work with it. Prime need is Accelerate can prepare an optimizer, but it converts it back to pytorch. We then need to wrap it again in OptimWrapper.

Example use case:

# some func to build a learner
learn = get_learner(...)
opt = accelerator.prepare_optimizer(learn.opt)
# Rewrap the optimizer
learn.opt = OptimWrapper(opt=opt)

To also have it be compatible with Accelerate, step takes in a closure param, but will raise an error if it's anything but None. This keeps it compatible with regular torch optimizers as well, though the implementation is not there (which @tmabraham is working on I believe)
cc @jph00

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB