fix

modeling_gptbert.py

@@ -439,8 +439,8 @@ class SelfAttention(nn.Module):
 
         else:
             # Standard attention path
-            query_length = hidden_layer.size(0)
-            key_length = hidden_layer.size(0)
+            query_length = query.size(1)
+            key_length = key.size(1)
 
             query = query.reshape(batch_size, query_length, self.num_attention_heads, self.d_qk).transpose(1, 2)
             key = key.reshape(batch_size, key_length, self.num_kv_heads, self.d_qk).transpose(1, 2)
@@ -451,7 +451,8 @@ class SelfAttention(nn.Module):
 
             if v1 is None:
                 v1 = value
-            value = (1 - self.lambdas[0]) * value + self.lambdas[0] * v1
+            else:
+                value = (1 - self.lambdas[0]) * value + self.lambdas[0] * v1
 
             # Apply rotary embeddings
             query = self.rope_embedding(query)
@@ -615,9 +616,9 @@ class GptBertModel(GptBertPreTrainedModel):
             padding_info = (indices, cu_seqlens, max_seqlen_in_batch)
         else:
             if len(attention_mask.size()) == 2:
+                attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+            elif len(attention_mask.size()) == 3:
                 attention_mask = attention_mask.unsqueeze(1)
-            if len(attention_mask.size()) != 3:
-                raise ValueError("Bare `attention_mask` med to eller tre dimensjoner støttes nå for SDPA.")
             padding_info = attention_mask
 
         static_embeddings = self.embedding(input_ids)