[TrainModel] Updated prototype for assignment