SWA implementation

[yuanqing-wang]

No description provided.

[yuanqing-wang]

Addresses https://forum.pyro.ai/t/using-stochastic-weight-averaging-swa-in-pyro/5290/2

[fritzo]

Hi @yuanqing-wang, I'd recommend taking a look at our CONTRIBUTING.md, and in particular:

• make format and make lint to work past linting errors
• add torchcontrib to the EXTRAS_REQUIRE list in setup.py (and ideally lower bound a version so we know SWA is available)
• add a test, probably in a new file tests/optim/test_swa.py
• probably add a pytest.importorskip("torchcontrib") to your test
• add a new automodule section to docs/source/pyro.optim.txt