[TrainModel] Updated classes for prototype

Added block traversal code