Module src.models.train_model