ott.neural.solvers.map_estimator.MapEstimator.train_map_estimator