Support SWA - Stochastic Weight Averaging optimizer

Is your feature request related to a problem? Please describe.

The PyTorch starting from version 1.6 support a set of tools to implement the Stochastic Weight Averaging (SWA) technique.

https://pytorch.org/blog/pytorch-1.6-now-includes-stochastic-weight-averaging/

They suggest that SWA can drastically improve the generalization performance and demonstrated to have a strong performance in several areas.

[1] Averaging Weights Leads to Wider Optima and Better Generalization; Pavel Izmailov, Dmitry Podoprikhin, Timur Garipov, Dmitry Vetrov, Andrew Gordon Wilson; Uncertainty in Artificial Intelligence (UAI), 2018.

Describe the solution you'd like

The SWA training it is not a drop-on replacement. It requires some very minor changes to the training loop. It would be nice to natively support a SWA training.