add use_buffers to swa_utils interface (#541)

Summary:
Pull Request resolved: #541

X-link: pytorch/pytorch#109078

As title, this already exists in swa_utils.py

Reviewed By: daniellepintz

Differential Revision: D49155243

fbshipit-source-id: df0e071247fa79d34eefae624bf46b896c62cb10

torchtnt/framework/auto_unit.py

@@ -54,10 +54,8 @@ def __call__(self, *args, **kwargs):
 
             if x.swa_params:
                 if not x.swa_params.avg_fn:
-                    # pyre-ignore: Unexpected keyword [28]
                     x.swa_model = AveragedModel(x.module, use_buffers=True)
                 else:
-                    # pyre-ignore: Unexpected keyword [28]
                     x.swa_model = AveragedModel(
                         x.module, avg_fn=x.swa_params.avg_fn, use_buffers=True
                     )