ott.core.neuraldual.NeuralDualSolver.create_train_state#

NeuralDualSolver.create_train_state(rng, model, optimizer, input)[source]#

Create initial TrainState.