Enable dynamo traced test_param_group_with_lrscheduler_goes_right_dir…

…ection (pytorch#124544) Pull Request resolved: pytorch#124544 Approved by: https://github.com/janeyx99 ghstack dependencies: pytorch#125825, pytorch#125826
mengpenghui · May 11, 2024 · b24ad7e · b24ad7e
1 parent e72ef4f
commit b24ad7e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 11 deletions.
diff --git a/test/test_optim.py b/test/test_optim.py
@@ -232,7 +232,9 @@ def test_param_group_with_lrscheduler_goes_right_direction(self, device, dtype,
             bias = Parameter(torch.randn((10), device=device, dtype=dtype))
             inpt = torch.randn(5, device=device, dtype=dtype)
 
-            optimizer = optim_cls([{"params": [weight]}, {"params": [bias], "lr": 0.01}])
+            # avoid endless recompiles by wrapping LR in a tensor if we're compiling
+            lr = torch.tensor(0.01) if torch._utils.is_compiling() else 0.01
+            optimizer = optim_cls([{"params": [weight]}, {"params": [bias], "lr": lr}])
             schedulers = [scheduler_c(optimizer) for scheduler_c in schedulers_c]
 
             def closure():

diff --git a/torch/testing/_internal/common_optimizers.py b/torch/testing/_internal/common_optimizers.py
@@ -1245,11 +1245,6 @@ def _get_optim_inputs_including_global_cliquey_kwargs(
                 "test_forloop_goes_right_direction",
                 active_if=lambda kwargs: not kwargs["contiguous"],
             ),
-            DecorateInfo(
-                skipIfTorchDynamo("initial_value is incorrect in dynamo, see #123202"),
-                "TestOptimRenewed",
-                "test_param_group_with_lrscheduler_goes_right_direction",
-            ),
             DecorateInfo(
                 skipIfTorchDynamo("Fails fix point assertion on 3.8, see #97811"),
                 "TestOptimRenewed",
@@ -1714,11 +1709,6 @@ def _get_optim_inputs_including_global_cliquey_kwargs(
             "cuda",
         ),
         skips=(
-            DecorateInfo(
-                skipIfTorchDynamo("initial_value is incorrect in dynamo, see #123202"),
-                "TestOptimRenewed",
-                "test_param_group_with_lrscheduler_goes_right_direction",
-            ),
             DecorateInfo(
                 skipIfTorchDynamo("Fails fix point assertion on 3.8, see #97811"),
                 "TestOptimRenewed",