Small fix in SDPA docstring codeblock (#109086)

Fix #109072 Pull Request resolved: #109086 Approved by: https://github.com/drisspg
pytorch · Sep 12, 2023 · abd83ce · abd83ce
1 parent 1b9b3a2
commit abd83ce
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch/nn/functional.py b/torch/nn/functional.py
@@ -4922,7 +4922,7 @@ def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.
 
         if attn_mask is not None:
             if attn_mask.dtype == torch.bool:
-                attn_mask.masked_fill_(attn_mask.logical_not(), float("-inf"))
+                attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
             else:
                 attn_bias += attn_mask
         attn_weight = query @ key.transpose(-2, -1) * scale_factor