Skip to content

GitLab

Unverified Commit 57d7594a authored by fxmarty's avatar fxmarty Committed by GitHub
Browse files

Fix mistral ONNX export (#31696) 

* use bitwise or

* why is the CI not triggered?
parent 93cd94b7
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
src/transformers/models/mistral/modeling_mistral.py
View file @ 57d7594a
...@@ -1089,8 +1089,9 @@ class MistralModel(MistralPreTrainedModel): ...@@ -1089,8 +1089,9 @@ class MistralModel(MistralPreTrainedModel):
exclude_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) exclude_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
if self.config.sliding_window is not None: if self.config.sliding_window is not None:
if not using_sliding_window_cache or sequence_length > self.config.sliding_window: if not using_sliding_window_cache or sequence_length > self.config.sliding_window:
exclude_mask |= torch.arange(target_length, device=device) <= ( exclude_mask.bitwise_or_(
cache_position.reshape(-1, 1) - self.config.sliding_window torch.arange(target_length, device=device)
<= (cache_position.reshape(-1, 1) - self.config.sliding_window)
) )
causal_mask *= exclude_mask causal_mask *= exclude_mask
causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment