Backout

Summary:
X-link: pytorch/FBGEMM#3803

X-link: facebookresearch/FBGEMM#887

Backout D68055168 as it seems to break pyper and causes S498612.

Differential Revision: D70996903

Author: spcyppt

torchrec/distributed/batched_embedding_kernel.py

@@ -190,7 +190,7 @@ def __init__(
         state: Dict[Any, Any] = {}
         param_group: Dict[str, Any] = {
             "params": [],
-            "lr": emb_module.optimizer_args.learning_rate_tensor,
+            "lr": emb_module.optimizer_args.learning_rate,
         }
 
         params: Dict[str, Union[torch.Tensor, ShardedTensor]] = {}
@@ -383,7 +383,7 @@ def get_optimizer_pointwise_shard_metadata_and_global_metadata(
         state: Dict[Any, Any] = {}
         param_group: Dict[str, Any] = {
             "params": [],
-            "lr": emb_module.optimizer_args.learning_rate_tensor,
+            "lr": emb_module.optimizer_args.learning_rate,
         }
 
         params: Dict[str, Union[torch.Tensor, ShardedTensor]] = {}

torchrec/modules/fused_embedding_modules.py

@@ -68,7 +68,7 @@ def __init__(  # noqa C901
         state: Dict[Any, Any] = {}
         param_group: Dict[str, Any] = {
             "params": [],
-            "lr": emb_module.optimizer_args.learning_rate_tensor,
+            "lr": emb_module.optimizer_args.learning_rate,
         }
 
         params: Dict[str, torch.Tensor] = {}