updated train functions