From b76e502a02b349c2a500253cd280db0aeab7aab3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 30 Jul 2024 23:07:47 -0700 Subject: [PATCH] Update gemma2.py --- unsloth/models/gemma2.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index a0880daefe..2191a99c07 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -130,11 +130,8 @@ def Gemma2Attention_fast_forward( past_key_value = (K, V) if use_cache else None # Only enable if the attention_mask is True - has_sliding_window = type(attention_mask) is bool and attention_mask is True - print(HAS_FLASH_ATTENTION_SOFTCAPPING) - print(has_sliding_window) - print(attention_mask) - if HAS_FLASH_ATTENTION_SOFTCAPPING and type(attention_mask) is bool: + has_sliding_window = type(causal_mask) is bool and causal_mask is True + if HAS_FLASH_ATTENTION_SOFTCAPPING and attention_mask is None: window = (-1, -1) if has_sliding_window: sw = getattr(self.config, "sliding_window", None)