ott.neural.networks.potentials.PotentialTrainState.create

classmethod PotentialTrainState.create(*, apply_fn, params, tx, **kwargs)

Creates a new instance with step=0 and initialized opt_state.