ott.neural.models.ICNN.create_train_state