TrainModel class has an incorrectly named variable - prevents loading a model checkpoint.

[kathyxchen]

self.max_step should be self.max_steps.

I will fix this in an upcoming PR.