fix buffer size

karpathy · Jul 23, 2024 · bc882c7 · bc882c7
1 parent 1e4a108
commit bc882c7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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) {