Back out "Fix dtype mismatch between weight and per_sample_weights" (#1791)

Summary:
Pull Request resolved: #1791

in inference tbe_input_combine op does not support fp16, we will need to revert to avoid any prod impact.

Original commit changeset: 969cb64c4af3

Original Phabricator Diff: D54526190

Reviewed By: peking2, houseroad

Differential Revision: D54884092

fbshipit-source-id: 728d89ae1ccb6fd219510c378d88c1a9ac20c880

Author: ZhengkaiZ

torchrec/modules/embedding_modules.py

@@ -162,7 +162,6 @@ def __init__(
         self.embedding_bags: nn.ModuleDict = nn.ModuleDict()
         self._embedding_bag_configs = tables
         self._lengths_per_embedding: List[int] = []
-        self._dtypes: List[int] = []
 
         table_names = set()
         for embedding_config in tables:
@@ -184,7 +183,6 @@ def __init__(
             )
             if device is None:
                 device = self.embedding_bags[embedding_config.name].weight.device
-            self._dtypes.append(embedding_config.data_type.value)
 
             if not embedding_config.feature_names:
                 embedding_config.feature_names = [embedding_config.name]
@@ -221,19 +219,10 @@ def forward(self, features: KeyedJaggedTensor) -> KeyedTensor:
         for i, embedding_bag in enumerate(self.embedding_bags.values()):
             for feature_name in self._feature_names[i]:
                 f = feature_dict[feature_name]
-                per_sample_weights: Optional[torch.Tensor] = None
-                if self._is_weighted:
-                    per_sample_weights = (
-                        f.weights().half()
-                        if self._dtypes[i] == DataType.FP16.value
-                        else f.weights()
-                    )
                 res = embedding_bag(
                     input=f.values(),
                     offsets=f.offsets(),
-                    per_sample_weights=(
-                        per_sample_weights if self._is_weighted else None
-                    ),
+                    per_sample_weights=f.weights() if self._is_weighted else None,
                 ).float()
                 pooled_embeddings.append(res)
         return KeyedTensor(