diff --git a/src/schema_compile_amend.c b/src/schema_compile_amend.c
index 9fff787..c45ee61 100644
--- a/src/schema_compile_amend.c
+++ b/src/schema_compile_amend.c
@@ -1887,9 +1887,11 @@
         lysc_update_path(ctx, NULL, NULL);
         LY_CHECK_GOTO(ret, cleanup);
 
-        /* augment was applied, remove it (index may have changed because other augments could have been applied) */
+        /* augment was applied, remove it (index and the whole set may have changed because other augments
+         * could have been applied) */
         ly_set_rm(&ctx->uses_augs, aug, NULL);
         lysc_augment_free(ctx->ctx, aug);
+        i = 0;
     }
 
     /* top-level augments */
@@ -1919,6 +1921,7 @@
         /* augment was applied, remove it */
         ly_set_rm(&ctx->augs, aug, NULL);
         lysc_augment_free(ctx->ctx, aug);
+        i = 0;
     }
 
 cleanup: