ott.solvers.nn.neuraldual.W2NeuralDual.setup

W2NeuralDual.setup(rng, neural_f, neural_g, dim_data, optimizer_f, optimizer_g)

Setup all components required to train the network.

Parameters:

• rng (PRNGKeyArray) –

• neural_f (ModelBase) –

• neural_g (ModelBase) –

• dim_data (int) –

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

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

Return type:

None