Fixing device issues (#1880)

Summary:
Pull Request resolved: #1880

Fixing device issues with sharding TorchRec inference modules on meta device

Reviewed By: IvanKobzarev

Differential Revision: D56020810

fbshipit-source-id: c14298130e702c1cbd67f1660f341178c15edbb4

Author: PaulZhang12

torchrec/distributed/embedding_lookup.py

@@ -917,7 +917,12 @@ def __init__(
                 MetaInferGroupedPooledEmbeddingsLookup(
                     grouped_configs=grouped_configs_per_rank[rank],
                     # syntax for torchscript
-                    device=torch.device(type=device_type, index=rank),
+                    # No rank for cpu
+                    device=(
+                        torch.device(type=device_type, index=rank)
+                        if device_type != "cpu"
+                        else torch.device(device_type)
+                    ),
                     fused_params=fused_params,
                 )
             )

torchrec/distributed/quant_embedding_kernel.py

@@ -163,7 +163,7 @@ def _unwrap_kjt(
 def _unwrap_kjt_for_cpu(
     features: KeyedJaggedTensor,
 ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
-    assert features.device().type == "cpu"
+    assert features.device().type == "cpu" or features.device().type == "meta"
     return features.values(), features.offsets(), features.weights_or_none()