Fix the predict bug that causes integration test to fail (sometimes)