ott.solvers.nn.models.ICNN.create_train_state

ICNN.create_train_state(rng, optimizer, input, **kwargs)

Create initial training state.

Parameters:

• rng (PRNGKeyArray) –

• optimizer (Union[Array, ndarray, bool_, number, Iterable[ArrayTree], Mapping[Any, ArrayTree]]) –

• input (Union[int, Tuple[int, ...]]) –

• kwargs (Any) –

Return type:

NeuralTrainState