Fixed bug in generating code for augment.