demonstrate

Birch-san · Sep 11, 2022 · d2d533d · d2d533d
1 parent 18bb5f8
commit d2d533d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ldm/modules/attention.py b/ldm/modules/attention.py
@@ -179,7 +179,7 @@ def forward(self, x, context=None, mask=None):
 
  q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
 
- sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+ sim = torch.matmul(q, k.transpose(1, 2)) * self.scale
  del q, k
 
  if exists(mask):
@@ -193,7 +193,7 @@ def forward(self, x, context=None, mask=None):
  attn = sim.softmax(dim=-1)
  del sim
 
- out = einsum('b i j, b j d -> b i d', attn, v)
+ out = torch.matmul(attn, v)
  del attn, v
  out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
  del h