[ML][Training] Updating code to native API changes