Minor change to code used in paper