[ML][Training] Model implementation