Update modular_starcoder2.py

huggingface · Dec 16, 2024 · 598b7bb · 598b7bb
1 parent 5e56d9c
commit 598b7bb
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/src/transformers/models/starcoder2/modular_starcoder2.py b/src/transformers/models/starcoder2/modular_starcoder2.py
@@ -132,14 +132,12 @@ def forward(
         if self.config._attn_implementation != "eager":
             attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
 
-        dropout_rate = 0.0 if not self.training else self.attention_dropout
-
         attn_output, attn_weights = attention_interface(
             self,
             query_states,
             key_states,
             value_states,
-            dropout=dropout_rate,
+            dropout=0.0 if not self.training else self.attention_dropout,
             **kwargs,
         )
 
@@ -283,7 +281,7 @@ class Starcoder2ForTokenClassification(LlamaForTokenClassification):
 __all__ = [
     "Starcoder2ForCausalLM",
     "Starcoder2Model",
-    "Starcoder2PreTrainedModel",
+    "Starcoder2PreTrainedModel", # noqa: F822
     "Starcoder2ForSequenceClassification",
     "Starcoder2ForTokenClassification",
 ]