Support output dtype for embeddings (#1744)

Summary:
Pull Request resolved: #1744

As title

Reviewed By: jiaqizhai

Differential Revision:
D54337769

Privacy Context Container: L1183554

fbshipit-source-id: d4abb56fd846e3da9897f2e295485e6bebb74584

torchrec/distributed/batched_embedding_kernel.py

@@ -666,6 +666,8 @@ def __init__(
         super().__init__(config, pg, device)
 
         weights_precision = data_type_to_sparse_type(config.data_type)
+        fused_params = config.fused_params or {}
+        output_dtype = fused_params.get("output_dtype", SparseType.FP32)
         self._emb_module: DenseTableBatchedEmbeddingBagsCodegen = (
             DenseTableBatchedEmbeddingBagsCodegen(
                 list(zip(self._local_rows, self._local_cols)),
@@ -675,6 +677,7 @@ def __init__(
                 or device.type == "cpu"
                 or not torch.cuda.is_available(),
                 weights_precision=weights_precision,
+                output_dtype=output_dtype,
             )
         )
         self._param_per_table: Dict[str, TableBatchedEmbeddingSlice] = dict(
@@ -958,6 +961,8 @@ def __init__(
         super().__init__(config, pg, device)
 
         weights_precision = data_type_to_sparse_type(config.data_type)
+        fused_params = config.fused_params or {}
+        output_dtype = fused_params.get("output_dtype", SparseType.FP32)
         self._emb_module: DenseTableBatchedEmbeddingBagsCodegen = (
             DenseTableBatchedEmbeddingBagsCodegen(
                 list(zip(self._local_rows, self._local_cols)),
@@ -967,6 +972,7 @@ def __init__(
                 or device.type == "cpu"
                 or not torch.cuda.is_available(),
                 weights_precision=weights_precision,
+                output_dtype=output_dtype,
             )
         )
         self._param_per_table: Dict[str, TableBatchedEmbeddingSlice] = dict(