Prediction bug fix