Add train class in train.rb