K!=Kv validation fix

ghstack-source-id: 5dbbd73acdeed410631da610c81453cdc5e1abcb Pull Request resolved: fairinternal/xformers#957 __original_commit__ = fairinternal/xformers@5af9cb3
facebookresearch · Dec 6, 2023 · af2f04b · af2f04b
1 parent 04baa4a
commit af2f04b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/xformers/ops/fmha/common.py b/xformers/ops/fmha/common.py
@@ -292,7 +292,7 @@ def not_supported_reasons(cls, d: Inputs) -> List[str]:
             Mq=d.query.shape[1],
             Mkv=d.key.shape[1],
             K=d.query.shape[-1],
-            Kv=d.query.shape[-1],
+            Kv=d.value.shape[-1],
         )
         device_type = d.query.device.type
         dtype = d.query.dtype