updates in logic for getting train classes