diff --git a/model/model.py b/model/model.py
index 010c8e3..6235e4e 100644
--- a/model/model.py
+++ b/model/model.py
@@ -627,7 +627,7 @@ class MiniMindLM(PreTrainedModel):
         # Specific layers for v path
         self.downsample_v_specific = nn.Sequential(
             nn.Conv1d(128*8, 128, kernel_size=1, padding='same'),
-            nn.Conv1d(128, 8, kernel_size=1, padding='same')
+            nn.Conv1d(128, self.params.knowledge_length, kernel_size=1, padding='same')
         )
 
         # Specific layers for q path