Fix wrong docstring arg name

facebookresearch · Aug 26, 2024 · 019bd37 · 019bd37
1 parent 00aba59
commit 019bd37
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/xformers/ops/fmha/__init__.py b/xformers/ops/fmha/__init__.py
@@ -677,7 +677,7 @@ def merge_attentions(
             either as a list of tensors of shapes [B, G, H, M] or [B, H, M]
             or as a single tensor of shape [num_chunks, B, G, H, M] or [num_chunks, B, H, M]
         write_lse: whether to output LSE
-        out_dype: dtype of attn_out
+        output_dtype: dtype of attn_out
 
     Returns:
         attn_out: [B, M, G, H, Kq] or [B, M, H, Kq]