Make this work with both train inputs.

llvm-svn: 12024
1 file changed