minor changes and factor out train script