diff --git a/examples/baby-llama/baby-llama-text.cpp b/examples/baby-llama/baby-llama-text.cpp
index 0c13f7fd4..4c3cd631b 100644
--- a/examples/baby-llama/baby-llama-text.cpp
+++ b/examples/baby-llama/baby-llama-text.cpp
@@ -68,7 +68,7 @@ struct ggml_tensor * randomize_tensor_normal(struct ggml_tensor * tensor, struct
             }
             break;
         case 2:
-            scale /= sqrtf(tensor->ne[0]*tensor->ne[1]);
+            scale /= sqrtf(tensor->ne[0]+tensor->ne[1]);
             for (int i1 = 0; i1 < tensor->ne[1]; i1++) {
                 for (int i0 = 0; i0 < tensor->ne[0]; i0++) {
                     float * dst = (float *) ((char *) tensor->data + i0*tensor->nb[0] + i1*tensor->nb[1]);
@@ -77,7 +77,7 @@ struct ggml_tensor * randomize_tensor_normal(struct ggml_tensor * tensor, struct
             }
             break;
         case 3:
-            scale /= sqrtf(tensor->ne[0]*tensor->ne[1]);
+            scale /= sqrtf(tensor->ne[0]+tensor->ne[1]);
             for (int i2 = 0; i2 < tensor->ne[2]; i2++) {
                 for (int i1 = 0; i1 < tensor->ne[1]; i1++) {
                     for (int i0 = 0; i0 < tensor->ne[0]; i0++) {
@@ -88,7 +88,7 @@ struct ggml_tensor * randomize_tensor_normal(struct ggml_tensor * tensor, struct
             }
             break;
         case 4:
-            scale /= sqrtf(tensor->ne[0]*tensor->ne[1]);
+            scale /= sqrtf(tensor->ne[0]+tensor->ne[1]);
             for (int i3 = 0; i3 < tensor->ne[3]; i3++) {
                 for (int i2 = 0; i2 < tensor->ne[2]; i2++) {
                     for (int i1 = 0; i1 < tensor->ne[1]; i1++) {