The OptimWrapper documentation suggests that using a pre-defined optimizer with Learner's opt_func argument should work:

Or if you already have an existing one, pass in only opt:

opt = torch.optim.SGD([tensor([1,2,3])], lr=1e-2)
opt_func = OptimWrapper(opt=opt)

however, following this example and passing opt_func to Learner.opt_func results in an error, as Learner.create_opt only accepts callable methods and not initialized objects.

This PR modifies Learner to accept a pre-initialized OptimWrapper to the opt_func argument by modifying Learner.create_opt to set a pre-initialized OptimWrapper to Learner.opt. To mimic the current behavior of always initializing an optimizer from scratch, create_opt will call clear_state on a pre-initialized OptimWrapper.

This allows a straightforward method of using custom parameter groups in a PyTorch style optimizer:

param_groups = custom_param_group_method(model)
opt = OptimWrapper(opt=optim.AdamW(param_groups, ...))
learn = Learner(dls, model, opt_func=opt, ...)

To prevent a pre-initialized optimizer's state from being cleared on the first fit, users can use the current behavior of passing the optimizer directly to Learner.opt. However, in this case manually calling Learner.fit(..., reset_opt=True) will cause an error. Instead of using reset_opt, users will need to redefine the optimizer.

This PR also updates some of the existing Learner docments to better match the docment style and more accurately describe the argument behavior.