Fixed update indices bug, fixed prob value in main