add missing annotation

pyro-ppl · Jun 6, 2018 · 76fa97c · 76fa97c
1 parent 4995fcf
commit 76fa97c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/infer/test_gradient.py b/tests/infer/test_gradient.py
@@ -29,12 +29,14 @@ def test_subsample_gradient(Elbo, reparameterized, subsample):
     precision = 0.06
     Normal = dist.Normal if reparameterized else fakes.NonreparameterizedNormal
 
+    @poutine.broadcast
     def model(subsample):
         with pyro.iarange("data", len(data), subsample_size, subsample) as ind:
             x = data[ind]
             z = pyro.sample("z", Normal(0, 1))
             pyro.sample("x", Normal(z, 1), obs=x)
 
+    @poutine.broadcast
     def guide(subsample):
         loc = pyro.param("loc", lambda: torch.zeros(len(data), requires_grad=True))
         scale = pyro.param("scale", lambda: torch.tensor([1.0], requires_grad=True))