ott.neural.networks.potentials.BasePotential.create_train_state# BasePotential.create_train_state(rng, optimizer, input, **kwargs)[source]# Create initial training state. Parameters: rng (Array) – optimizer (Union[Array, ndarray, bool_, number, Iterable[ArrayTree], Mapping[Any, ArrayTree]]) – input (Union[int, Tuple[int, ...]]) – kwargs (Any) – Return type: PotentialTrainState