Add explicit init() to BasePotential and NeuralPosterior

sbi/inference/posteriors/base_posterior.py

@@ -79,6 +79,14 @@ def __init__(
         # as default x.
         self._x = self.potential_fn.return_x_o()
 
+        self._initialized = False
+
+    def init(self):
+        """Initialize the underlying potential before sampling."""
+        self.potential_fn.init()
+        self._initialized = True
+        return self
+
     def potential(
         self, theta: Tensor, x: Optional[Tensor] = None, track_gradients: bool = False
     ) -> Tensor:
@@ -93,6 +101,8 @@ def potential(
                 This can be helpful for e.g. sensitivity analysis, but increases memory
                 consumption.
         """
+        if not self._initialized:
+            self.init()
         self.potential_fn.set_x(self._x_else_default_x(x))
 
         theta = ensure_theta_batched(torch.as_tensor(theta))

sbi/inference/potentials/base_potential.py

@@ -30,6 +30,12 @@ def __init__(
         self.device = device
         self.prior = prior
         self.set_x(x_o)
+        self._initialized = False
+
+    def init(self):
+        """Initialize potential (can be overridden by subclasses)."""
+        self._initialized = True
+        return self
 
     @abstractmethod
     def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor:
@@ -136,5 +142,7 @@ def __call__(self, theta, track_gradients: bool = True):
 
         Note, x_o is re-used from the initialization of the potential function.
         """
+        if not self._initialized:
+            self.init()
         with torch.set_grad_enabled(track_gradients):
             return self.potential_fn(theta, self.x_o)