error in loss augmented inference fixed
minor fixes