fix pyre errors (#784)

Summary: Pull Request resolved: #784

Reviewed By: anshulverma

Differential Revision: D56144363

fbshipit-source-id: 052efee08a0528416477b132990c2453450d2f35

Author: galrotem

torchtnt/utils/swa.py

@@ -10,8 +10,6 @@
 
 import torch
 
-# TODO: torch/optim/swa_utils.pyi needs to be updated
-# pyre-ignore Undefined import [21]: Could not find a name `get_ema_multi_avg_fn` defined in module `torch.optim.swa_utils`.
 from torch.optim.swa_utils import (
     AveragedModel as PyTorchAveragedModel,
     get_ema_multi_avg_fn,
@@ -56,12 +54,8 @@ def __init__(
             if ema_decay < 0.0 or ema_decay > 1.0:
                 raise ValueError(f"Decay must be between 0 and 1, got {ema_decay}")
 
-            # TODO: torch/optim/swa_utils.pyi needs to be updated
-            # pyre-ignore Undefined attribute [16]: Module `torch.optim.swa_utils` has no attribute `get_ema_multi_avg_fn`.
             multi_avg_fn = get_ema_multi_avg_fn(ema_decay)
         elif averaging_method == "swa":
-            # TODO: torch/optim/swa_utils.pyi needs to be updated
-            # pyre-ignore Undefined attribute [16]: Module `torch.optim.swa_utils` has no attribute `get_swa_multi_avg_fn`.
             multi_avg_fn = get_swa_multi_avg_fn()
 
             if use_lit:
@@ -89,8 +83,6 @@ def __init__(
         else:
             # use default init implementation
 
-            # TODO: torch/optim/swa_utils.pyi needs to be updated
-            # pyre-ignore Unexpected keyword [28]
             super().__init__(
                 model,
                 device=device,
@@ -105,7 +97,5 @@ def update_parameters(self, model: torch.nn.Module) -> None:
                 self._ema_decay, (1 + self._num_updates) / (10 + self._num_updates)
             )
 
-            # TODO: torch/optim/swa_utils.pyi needs to be updated
-            # pyre-ignore Undefined attribute [16]: Module `torch.optim.swa_utils` has no attribute `get_ema_multi_avg_fn`.
             self.multi_avg_fn = get_ema_multi_avg_fn(decay)
         super().update_parameters(model)