added new train function
small changes