Fix incorrect statement ordering in augment