Fix pipeline test/train bug