Fix the sync point caused by iter_cpu.item() (pytorch#489)

Summary:
X-link: pytorch#3401

Pull Request resolved: facebookresearch/FBGEMM#489

Although `self.iter_cpu` is created on CPU. It might be transferred to
GPU by the user. So, we need to transfer it to CPU explicitly. This
should be done only once.

Reviewed By: csmiler

Differential Revision: D66311970

fbshipit-source-id: 4ced9b28f2e6c69fc12ccdc7e88404caf8276627

Author: sryap

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -1824,6 +1824,12 @@ def forward(  # noqa: C901
             #  `Union[Module, Tensor]`.
             placements=self.momentum2_placements,
         )
+
+        # Although self.iter_cpu is created on CPU. It might be transferred to
+        # GPU by the user. So, we need to transfer it to CPU explicitly. This
+        # should be done only once.
+        self.iter_cpu = self.iter_cpu.cpu()
+
         # Sync with loaded state
         if (
             not is_torchdynamo_compiling()