Upload modeling_vit.py

modeling_vit.py

@@ -325,9 +325,9 @@ class ViTSdpaSelfAttention(ViTSelfAttention):
             query_layer,
             key_layer,
             value_layer,
-            head_mask,
+            dropout_p=self.attention_probs_dropout_prob if self.training else 0.0,
+            attn_mask=head_mask,
             softmax_fn = self.softmax_fn,
-            self.attention_probs_dropout_prob if self.training else 0.0,
             is_causal=False,
             scale=None,
         )