Tweak parameters in inference test

pyro-ppl · Feb 23, 2021 · 10309f2 · 10309f2
1 parent 6f4cad0
commit 10309f2
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/contrib/autoguide/test_inference.py b/tests/contrib/autoguide/test_inference.py
@@ -80,7 +80,7 @@ def do_test_auto(self, N, reparameterized, n_steps):
                                         AutoLowRankMultivariateNormal, AutoLaplaceApproximation])
 @pytest.mark.parametrize('Elbo', [Trace_ELBO, TraceMeanField_ELBO])
 def test_auto_diagonal_gaussians(auto_class, Elbo):
-    n_steps = 3501 if auto_class == AutoDiagonalNormal else 6001
+    n_steps = 3001
 
     def model():
         pyro.sample("x", dist.Normal(-0.2, 1.2))
@@ -90,7 +90,8 @@ def model():
         guide = auto_class(model, rank=1)
     else:
         guide = auto_class(model)
-    adam = optim.Adam({"lr": .001, "betas": (0.95, 0.999)})
+    adam = optim.ClippedAdam({"lr": .01, "betas": (0.95, 0.999),
+                              "lrd": 0.1 ** (1 / n_steps)})
     svi = SVI(model, guide, adam, loss=Elbo())
 
     for k in range(n_steps):