Add annotations to functions in stl/lightning/utilities/ema_state.py

Summary:
Add annotations to the following functions:
- def __init__(...)
- def _get_model_state_iterator(...)

Reviewed By: kandluis

Differential Revision: D27745227

fbshipit-source-id: 60b15ce87bf7e61ed17979f34c26986491bfdb4e

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py

@@ -709,8 +709,6 @@ def split_embedding_weights(self) -> List[Tensor]:
 
     @torch.jit.ignore
     def get_optimizer_buffer(self, state: str) -> torch.Tensor:
-        # pyre-fixme[16]: `SplitTableBatchedEmbeddingBagsCodegen` has no
-        #  attribute `named_buffers`.
         for name, buffer in self.named_buffers():
             if name == state:
                 return buffer