diff --git a/src/schema_compile_amend.c b/src/schema_compile_amend.c
index a751ec4..4e59b02 100644
--- a/src/schema_compile_amend.c
+++ b/src/schema_compile_amend.c
@@ -1852,14 +1852,9 @@
     ly_bool enabled, child_unres_disabled = 0;
     uint32_t opt_prev = ctx->compile_opts;
 
+    assert(target->nodetype & (LYS_CONTAINER | LYS_LIST | LYS_CHOICE | LYS_CASE | LYS_INPUT | LYS_OUTPUT | LYS_NOTIF));
+
     /* nodetype checks */
-    if (!(target->nodetype & (LYS_CONTAINER | LYS_LIST | LYS_CHOICE | LYS_CASE | LYS_INPUT | LYS_OUTPUT | LYS_NOTIF))) {
-        LOGVAL(ctx->ctx, LYVE_REFERENCE,
-                "Augment's %s-schema-nodeid \"%s\" refers to a %s node which is not an allowed augment's target.",
-                aug_p->nodeid[0] == '/' ? "absolute" : "descendant", aug_p->nodeid, lys_nodetype2str(target->nodetype));
-        rc = LY_EVALID;
-        goto cleanup;
-    }
     if (aug_p->actions && !lysc_node_actions_p(target)) {
         LOGVAL(ctx->ctx, LYVE_REFERENCE,
                 "Invalid augment of %s node which is not allowed to contain RPC/action node \"%s\".",
diff --git a/src/schema_compile_node.c b/src/schema_compile_node.c
index 02a1e51..30671d1 100644
--- a/src/schema_compile_node.c
+++ b/src/schema_compile_node.c
@@ -2563,9 +2563,6 @@
         LY_CHECK_GOTO(ret, cleanup);
     }
 
-    /* connect any augments */
-    LY_CHECK_GOTO(ret = lys_compile_node_augments(ctx, node), cleanup);
-
     /* nodetype-specific part */
     LY_CHECK_GOTO(ret = node_compile_spec(ctx, pnode, node), cleanup);
 
@@ -2620,6 +2617,9 @@
         LY_CHECK_GOTO(ret = lys_compile_node(ctx, child_p, node, 0, NULL), done);
     }
 
+    /* connect any augments */
+    LY_CHECK_GOTO(ret = lys_compile_node_augments(ctx, node), done);
+
     ctx->compile_opts = prev_options;
 
 done:
@@ -2712,6 +2712,9 @@
         LY_CHECK_GOTO(ret, done);
     }
 
+    /* connect any augments */
+    LY_CHECK_GOTO(ret = lys_compile_node_augments(ctx, node), done);
+
 done:
     return ret;
 }
@@ -2753,6 +2756,9 @@
     ret = lysc_unres_must_add(ctx, node, pnode);
     LY_CHECK_GOTO(ret, done);
 
+    /* connect any augments */
+    LY_CHECK_GOTO(ret = lys_compile_node_augments(ctx, node), done);
+
     LY_LIST_FOR((struct lysp_node *)cont_p->actions, child_p) {
         ret = lys_compile_node(ctx, child_p, node, 0, NULL);
         LY_CHECK_GOTO(ret, done);
@@ -3344,6 +3350,9 @@
         lysc_update_path(ctx, NULL, NULL);
     }
 
+    /* connect any augments */
+    LY_CHECK_GOTO(ret = lys_compile_node_augments(ctx, node), done);
+
     /* uniques */
     if (list_p->uniques) {
         LY_CHECK_RET(lys_compile_node_list_unique(ctx, list_p->uniques, list));
@@ -3512,14 +3521,18 @@
     assert(node->nodetype == LYS_CHOICE);
 
     LY_LIST_FOR(ch_p->child, child_p) {
-        LY_CHECK_RET(lys_compile_node_choice_child(ctx, child_p, node, NULL));
+        LY_CHECK_GOTO(ret = lys_compile_node_choice_child(ctx, child_p, node, NULL), done);
     }
 
+    /* connect any augments */
+    LY_CHECK_GOTO(ret = lys_compile_node_augments(ctx, node), done);
+
     /* default branch */
     if (ch_p->dflt.str) {
-        LY_CHECK_RET(lys_compile_node_choice_dflt(ctx, &ch_p->dflt, ch));
+        LY_CHECK_GOTO(ret = lys_compile_node_choice_dflt(ctx, &ch_p->dflt, ch), done);
     }
 
+done:
     return ret;
 }
 
@@ -3569,6 +3582,7 @@
 static LY_ERR
 lys_compile_node_case(struct lysc_ctx *ctx, struct lysp_node *pnode, struct lysc_node *node)
 {
+    LY_ERR ret = LY_SUCCESS;
     struct lysp_node *child_p;
     struct lysp_node_case *cs_p = (struct lysp_node_case *)pnode;
 
@@ -3577,13 +3591,17 @@
     } else if (pnode->nodetype == LYS_CASE) {
         /* explicit parent case */
         LY_LIST_FOR(cs_p->child, child_p) {
-            LY_CHECK_RET(lys_compile_node(ctx, child_p, node, 0, NULL));
+            LY_CHECK_GOTO(ret = lys_compile_node(ctx, child_p, node, 0, NULL), done);
         }
     } else {
         LOGINT_RET(ctx->ctx);
     }
 
-    return LY_SUCCESS;
+    /* connect any augments */
+    LY_CHECK_GOTO(ret = lys_compile_node_augments(ctx, node), done);
+
+done:
+    return ret;
 }
 
 void