Update blocks.py

9180c729 · mashun1 · c1676293 · 9180c729
Commit 9180c729 authored Aug 30, 2024 by mashun1
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

opensora/models/layers/blocks.py opensora/models/layers/blocks.py +1 -1

No files found.
--- a/opensora/models/layers/blocks.py
+++ b/opensora/models/layers/blocks.py
@@ -282,7 +282,7 @@ class MultiHeadCrossAttention(nn.Module):
        attn_bias = None
        if mask is not None:
            attn_bias = xformers.ops.fmha.BlockDiagonalMask.from_seqlens([N] * B, mask)
-        x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, attn_bias=attn_bias)
+        x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, attn_bias=attn_bias, op=xformers.ops.fmha.MemoryEfficientAttentionFlashAttentionOp)

        x = x.view(B, -1, C)
        x = self.proj(x)