Fixed small bug in train example script