Update modeling_auristream.py

modeling_auristream.py

@@ -445,6 +445,9 @@ class CausalSelfAttention(nn.Module):
                 is_causal=True)
         else:
             # manual implementation of attention
+            q = q.transpose(1, 2)
+            k = k.transpose(1, 2)
+            v = v.transpose(1, 2)
             att = torch.einsum('bnsh,bnkh->bnsk', q, k) * (1.0 / math.sqrt(k.size(-1)))
             mask = torch.triu(torch.ones(T, T), diagonal=1).to(dtype=torch.bool).to(x.device)
             mask = mask.view(1, 1, T, T)