diff --git a/examples/train-text-from-scratch/train-text-from-scratch.cpp b/examples/train-text-from-scratch/train-text-from-scratch.cpp
index a4a6b05b1..c76a80c75 100644
--- a/examples/train-text-from-scratch/train-text-from-scratch.cpp
+++ b/examples/train-text-from-scratch/train-text-from-scratch.cpp
@@ -1403,8 +1403,8 @@ void graph_set_leafs_grads(struct ggml_cgraph * g) {
         }
     }
     for (int i=n_nodes; i < g->n_nodes; ++i) {
-        g->nodes[n_nodes] = NULL;
-        g->grads[n_nodes] = NULL;
+        g->nodes[i] = NULL;
+        g->grads[i] = NULL;
     }
     g->n_nodes = n_nodes;
 }