fix optimizer config to apply correctly in tests (pytorch#2436)

Summary: Pull Request resolved: pytorch#2436 tsia, same as D63468297 Differential Revision: D63555376
iamzainhuda · Sep 30, 2024 · 4ec4a5e · 4ec4a5e
1 parent db14fef
commit 4ec4a5e
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 8 deletions.
diff --git a/torchrec/distributed/tests/test_model_parallel_hierarchical.py b/torchrec/distributed/tests/test_model_parallel_hierarchical.py
@@ -73,7 +73,7 @@ class ModelParallelHierarchicalTest(ModelParallelTestShared):
             [
                 None,
                 {
-                    "embeddingbags": (torch.optim.SGD, {"lr": 0.01}),
+                    "embedding_bags": (torch.optim.SGD, {"lr": 0.01}),
                     "embeddings": (torch.optim.SGD, {"lr": 0.2}),
                 },
             ]
@@ -162,7 +162,7 @@ def test_sharding_nccl_twrw(
             [
                 None,
                 {
-                    "embeddingbags": (torch.optim.SGD, {"lr": 0.01}),
+                    "embedding_bags": (torch.optim.SGD, {"lr": 0.01}),
                     "embeddings": (torch.optim.SGD, {"lr": 0.2}),
                 },
             ]
@@ -286,7 +286,7 @@ def test_sharding_empty_rank(
             [
                 None,
                 {
-                    "embeddingbags": (torch.optim.SGD, {"lr": 0.01}),
+                    "embedding_bags": (torch.optim.SGD, {"lr": 0.01}),
                     "embeddings": (torch.optim.SGD, {"lr": 0.2}),
                 },
             ]
@@ -355,7 +355,7 @@ def test_embedding_tower_nccl(
             [
                 None,
                 {
-                    "embeddingbags": (torch.optim.SGD, {"lr": 0.01}),
+                    "embedding_bags": (torch.optim.SGD, {"lr": 0.01}),
                     "embeddings": (torch.optim.SGD, {"lr": 0.2}),
                 },
             ]

diff --git a/torchrec/distributed/tests/test_sequence_model_parallel.py b/torchrec/distributed/tests/test_sequence_model_parallel.py
@@ -57,7 +57,7 @@ class SequenceModelParallelTest(MultiProcessTestBase):
             [
                 None,
                 {
-                    "embeddingbags": (torch.optim.SGD, {"lr": 0.01}),
+                    "embedding_bags": (torch.optim.SGD, {"lr": 0.01}),
                     "embeddings": (torch.optim.SGD, {"lr": 0.2}),
                 },
             ]
@@ -150,7 +150,7 @@ def test_sharding_nccl_dp(
             [
                 None,
                 {
-                    "embeddingbags": (torch.optim.SGD, {"lr": 0.01}),
+                    "embedding_bags": (torch.optim.SGD, {"lr": 0.01}),
                     "embeddings": (torch.optim.SGD, {"lr": 0.2}),
                 },
             ]
@@ -203,7 +203,7 @@ def test_sharding_nccl_tw(
             [
                 None,
                 {
-                    "embeddingbags": (torch.optim.SGD, {"lr": 0.01}),
+                    "embedding_bags": (torch.optim.SGD, {"lr": 0.01}),
                     "embeddings": (torch.optim.SGD, {"lr": 0.2}),
                 },
             ]

diff --git a/torchrec/distributed/tests/test_sequence_model_parallel_hierarchical.py b/torchrec/distributed/tests/test_sequence_model_parallel_hierarchical.py
@@ -62,7 +62,7 @@ class SequenceModelParallelHierarchicalTest(MultiProcessTestBase):
             [
                 None,
                 {
-                    "embeddingbags": (torch.optim.SGD, {"lr": 0.01}),
+                    "embedding_bags": (torch.optim.SGD, {"lr": 0.01}),
                     "embeddings": (torch.optim.SGD, {"lr": 0.2}),
                 },
             ]
-Original file line number
+Diff line change
@@ Expand Up @@
                 [
                     None,
                     {
-                        "embeddingbags": (torch.optim.SGD, {"lr": 0.01}),
+                        "embedding_bags": (torch.optim.SGD, {"lr": 0.01}),
                         "embeddings": (torch.optim.SGD, {"lr": 0.2}),
                     },
                 ]
@@ Expand Down @@