Update part of hw2 (predict not correct)