ott.neural.networks.potentials.PotentialTrainState.potential_value_fn

PotentialTrainState.potential_value_fn: Callable[[FrozenDict[str, Array], Optional[Callable[[Array], Array]]], Callable[[Array], Array]]