Pyre fix (pytorch#2605)

Summary: Pull Request resolved: pytorch#2605

Reviewed By: sarckk

Differential Revision: D66717161

fbshipit-source-id: c347b7330ec7f7426c1162d84910387ade0997f0

torchrec/distributed/model_parallel.py

@@ -746,6 +746,7 @@ def sync(self, include_optimizer_state: bool = True) -> None:
         all_weights = [
             w
             for emb_kernel in self._modules_to_sync
+            # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
             for w in emb_kernel.split_embedding_weights()
         ]
         handle = self._replica_pg.allreduce_coalesced(all_weights, opts=opts)
@@ -755,6 +756,7 @@ def sync(self, include_optimizer_state: bool = True) -> None:
             # Sync accumulated square of grad of local optimizer shards
             optim_list = []
             for emb_kernel in self._modules_to_sync:
+                # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
                 all_optimizer_states = emb_kernel.get_optimizer_state()
                 momentum1 = [optim["sum"] for optim in all_optimizer_states]
                 optim_list.extend(momentum1)
@@ -864,6 +866,8 @@ def _find_sharded_modules(
             if isinstance(module, SplitTableBatchedEmbeddingBagsCodegen):
                 sharded_modules.append(module)
             if hasattr(module, "_lookups"):
+                # pyre-fixme[29]: `Union[(self: Tensor) -> Any, Module, Tensor]` is
+                #  not a function.
                 for lookup in module._lookups:
                     _find_sharded_modules(lookup)
                 return