Improve predict