Added train class to train.rb