Fix Pyre (pytorch#2960)

Summary: Pull Request resolved: pytorch#2960

Reviewed By: aliafzal

Differential Revision: D74418088

fbshipit-source-id: 6ce96e7833d162f65598354edcab5c71c1f1edca

Author: aporialiao

torchrec/distributed/tests/test_pt2_multiprocess.py

@@ -251,7 +251,6 @@ def _test_compile_rank_fn(
         num_float_features: int = 8
         num_weighted_features: int = 1
 
-        # pyre-ignore
         device: torch.Device = torch.device("cuda")
         pg: Optional[dist.ProcessGroup] = ctx.pg
         assert pg is not None
@@ -336,7 +335,7 @@ def _dmp(m: torch.nn.Module) -> DistributedModelParallel:
                 env=ShardingEnv.from_process_group(pg),
                 plan=plan,
                 sharders=sharders,
-                device=device,
+                device=device,  # pyre-ignore
                 init_data_parallel=False,
             )
 
@@ -601,7 +600,7 @@ def _dmp(m: torch.nn.Module) -> DistributedModelParallel:  # pyre-ignore
             env=ShardingEnv(world_size, rank, pg),
             plan=plan,
             sharders=sharders,
-            device=device,
+            device=device,  # pyre-ignore
             init_data_parallel=False,
         )