Adding train class
Adding train class