Fix MNLE test

sbi-dev · Dec 4, 2023 · 83a8d8f · 83a8d8f
1 parent 618a262
commit 83a8d8f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/mnle_test.py b/tests/mnle_test.py
@@ -243,7 +243,7 @@ def sim_wrapper(theta):
 
     # MNLE
     trainer = MNLE(proposal)
-    estimator = trainer.append_simulations(theta, x).train(max_num_epochs=1)
+    estimator = trainer.append_simulations(theta, x).train()
 
     potential_fn = MixedLikelihoodBasedPotential(estimator, proposal, x_o)
 
@@ -276,6 +276,7 @@ def sim_wrapper(theta):
         potential_fn=conditioned_potential_fn,
         theta_transform=prior_transform,
         proposal=prior,
+        **mcmc_kwargs,
     )
     cond_samples = mcmc_posterior.sample((num_samples,), x=x_o)