fix typo in batch size

danielward27 · Aug 26, 2023 · fdb3c38 · fdb3c38
1 parent 78c9350
commit fdb3c38
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flowjax/tasks.py b/flowjax/tasks.py
@@ -65,7 +65,7 @@ def sample_reference_posterior(
             key, subkey = jr.split(key)
 
             candidates = jax.jit(jax.vmap(self.simulator, in_axes=[0, None]))(
-                jr.split(subkey, 500), observation
+                jr.split(subkey, num_samples), observation
             )
             in_prior_support = self.prior.log_prob(candidates) != -jnp.inf
             sample_counter += in_prior_support.sum()