Fix prediction bug