fix optimwrapper to work with param_groups (closes #2829)

[tmabraham]

Currently, OptimWrapper does not support param_groups, preventing use of model freezing/unfreezing and discriminative learning rate.

This PR incorporates @muellerzr's fix into OptimWrapper. In doing so, the usage of OptimWrapper is slightly changed. For example, it changes usage from:

def opt_func(params, lr, **kwargs): return OptimWrapper(torch.optim.SGD(params, lr=lr))

to

def opt_func(params, **kwargs): return OptimWrapper(torch.optim.SGD, params, **kwargs)

This PR will also solve issue #2829 as well.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[tmabraham]

I just realized that by switching the arguments of OptimWrapper, I can make it even easier to use:
opt_func = partial(OptimWrapper, opt=torch.optim.SGD)

I have added this change as a commit, but @jph00 please let me know if you are fine with this. If you think it's too much change to the OptimWrapper usage, I will revert back to the previously agreed upon changes...

[jph00]

Looking good! Minor suggestions

[tmabraham]

@jph00 I made changes based on your suggestions, and also added the doc changes from #3226

[jph00]

@tmabraham any chance you could look to see why one notebook is timing out in tests in CI?

[tmabraham]

@jph00 fixed! surprised this wasn't an issue before but one of the cells was not marked slow...

[jph00]

Thank you! :)