From 3be14462eb2df99fd4dd190ae7745700f6372013 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 24 Dec 2024 00:11:17 -0800 Subject: [PATCH] Update cohere.py --- unsloth/models/cohere.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py index cbbebeec5c..1610949f64 100644 --- a/unsloth/models/cohere.py +++ b/unsloth/models/cohere.py @@ -68,7 +68,7 @@ pass def CohereAttention_fast_forward( self, hidden_states: torch.Tensor, - causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None, + causal_mask: Optional[BlockDiagonalCausalMask] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Tuple[torch.Tensor]] = None, @@ -183,7 +183,7 @@ pass def CohereDecoderLayer_fast_forward( self, hidden_states: torch.Tensor, - causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None, + causal_mask: Optional[BlockDiagonalCausalMask] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Tuple[torch.Tensor]] = None,