Quick fixes before the ci run.
Merge pull request #149 from roydude/v1-train

Quick fixes before the ci run.