ott.neural.solvers.neuraldual.BaseW2NeuralDual.create_train_state