remove .int() for cpu indices and values (#1590)

Summary:
Pull Request resolved: #1590

CPU DI is serving large model with big embedding tables (2TB), the value and indices would overflow with .int() conversion. Remove .int() just for CPU

Reviewed By: zyan0, tissue3

Differential Revision: D52225777

fbshipit-source-id: 0bf7973a91a7b7daed6eaed3a55bb8dca25fcdef

torchrec/distributed/quant_embedding_kernel.py

@@ -118,7 +118,14 @@ def _quantize_weight(
 def _unwrap_kjt(
     features: KeyedJaggedTensor,
 ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
-    return features.values().int(), features.offsets().int(), features.weights_or_none()
+    if features.device().type == "cuda":
+        return (
+            features.values().int(),
+            features.offsets().int(),
+            features.weights_or_none(),
+        )
+    else:
+        return features.values(), features.offsets(), features.weights_or_none()
 
 
 class QuantBatchedEmbeddingBag(