Represent query_length in a different way to solve jit issue (#25164)

Fix jit trace
huggingface · Jul 28, 2023 · d23d2c2 · d23d2c2
1 parent 2a78720
commit d23d2c2
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/src/transformers/models/mpt/modeling_mpt.py b/src/transformers/models/mpt/modeling_mpt.py
@@ -154,9 +154,7 @@ def forward(
 
         attention_scores = torch.matmul(query_states, key_states.transpose(-1, -2)) * self.softmax_scale
 
-        query_length = seq_length
-        if past_key_value is not None:
-            query_length += past_key_value[0].shape[2]
+        query_length = seq_length if past_key_value is None else seq_length + past_key_value[0].shape[2]
 
         if position_bias is not None:
             if len(position_bias.shape) != 3: