From bc882c7005159a202ac3454995071fa562a7537e Mon Sep 17 00:00:00 2001
From: Erik Schultheis <erik.schultheis@aalto.fi>
Date: Tue, 23 Jul 2024 23:02:21 +0200
Subject: [PATCH] fix buffer size

---
 train_gpt2.cu | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train_gpt2.cu b/train_gpt2.cu
index 1e8db96ea..50b2cb7fa 100644
--- a/train_gpt2.cu
+++ b/train_gpt2.cu
@@ -376,7 +376,7 @@ void gpt2_allocate_state(GPT2 *model, int B, int T) {
         size_t param_elements[NUM_PARAMETER_TENSORS];
         size_t param_sizeof[NUM_PARAMETER_TENSORS];
         GPT2Config wave_config = model->config;
-        wave_config.num_layers = 1;
+        wave_config.num_layers = 2;
         fill_in_parameter_sizes(param_elements, param_sizeof, wave_config);
         size_t alloc_bytes = 0;
         for(int i = 0; i < NUM_PARAMETER_TENSORS; ++i) {