setu.save_estimator#

setu.save_estimator(estimator, path)[source]#

Save a trained NLE, MixedNLE, or DiscreteNLE to disk.

Uses equinox serialization to save all array leaves of the estimator pytree.

Parameters:

estimator (NLE | MixedNLE | DiscreteNLE) – Trained NLE or MixedNLE instance.
path (str | Path) – File path to save to (e.g., “model.eqx”).

Return type:

None

Example

>>> nle = create_nle(x_dim=3, theta_dim=3, key=key)
>>> result = fit_nle(nle, dataset, key=key)
>>> save_estimator(result.nle, "trained_nle.eqx")