Merge pull request #99 from jialin-zhao/jialin-zhao-xformers-inputs-revised

Update latte.py about the xformers input dims Former-commit-id: 421b49c351124170ab476e17cba6e45bd7cf706e
2024-03-09 09:58:51 +08:00
parent c42a86310d cd2ffd7229
commit 38c2c2ccaf
1 changed files with 2 additions and 0 deletions
@@ -71,6 +71,8 @@ class Attention(nn.Module):
        q, k, v = qkv.unbind(0)   # make torchscript happy (cannot use tensor as tuple)
        
        if self.attention_mode == 'xformers': # cause loss nan while using with amp
+            assert q.ndim == 4
+            q, k, v = q.transpose(1,2), k.transpose(1,2), v.transpose(1,2)
            x = xformers.ops.memory_efficient_attention(q, k, v).reshape(B, N, C)

        elif self.attention_mode == 'flash':