Should fix OOMs in CI

ghstack-source-id: 0e543590f9fca47bd4ff1ab81dc60a37292d28c6 Pull Request resolved: #488
facebookresearch · Oct 21, 2022 · bcc40dc · bcc40dc
1 parent 0da65ca
commit bcc40dc
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/xformers/ops/memory_efficient_attention.py b/xformers/ops/memory_efficient_attention.py
@@ -633,6 +633,7 @@ def _flash_attn_backward(
 class MemoryEfficientAttentionCutlassFwdFlashBwOp(MemoryEfficientAttentionCutlassOp):
     FW_OP = MemoryEfficientAttentionCutlassOp
     BW_OP = MemoryEfficientAttentionFlashAttentionOp
+    SUPPORTED_DTYPES = BW_OP.SUPPORTED_DTYPES.intersection(FW_OP.SUPPORTED_DTYPES)
     NAME = "fctls_bflsh"
 
     @classmethod