Code update for Matei's suggestions