update infer() method

sbi/inference/base.py

@@ -29,10 +29,17 @@
 
 
 def infer(
-    simulator: Callable, prior, method: str, num_simulations: int, num_workers: int = 1
+    simulator: Callable,
+    prior: Any,
+    method: str,
+    num_simulations: int,
+    num_workers: int = 1,
 ) -> NeuralPosterior:
     r"""
-    Return posterior distribution by running simulation-based inference.
+    Runs simulation-based inference.
+
+    After running this, you will have to run `inference.build_posterior(prior, x_o)` to
+    obtain the posterior.
 
     This function provides a simple interface to run sbi. Inference is run for a single
     round and hence the returned posterior $p(\theta|x)$ can be sampled and evaluated
@@ -69,17 +76,16 @@ def infer(
 
     simulator, prior = prepare_for_sbi(simulator, prior)
 
-    inference = method_fun(prior)
+    inference = method_fun()
     theta, x = simulate_for_sbi(
         simulator=simulator,
         proposal=prior,
         num_simulations=num_simulations,
         num_workers=num_workers,
     )
     _ = inference.append_simulations(theta, x).train()
-    posterior = inference.build_posterior()
 
-    return posterior
+    return inference
 
 
 class NeuralInference(ABC):