some reference-freeing I missed from @neonsecret's optimized attention …

…CompVis#177 plus some extras from me (del v, del h)
Birch-san · Sep 11, 2022 · 18bb5f8 · 18bb5f8
1 parent 91d29c2
commit 18bb5f8
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/ldm/modules/attention.py b/ldm/modules/attention.py
@@ -191,9 +191,12 @@ def forward(self, x, context=None, mask=None):
 
         # attention, what we cannot get enough of
         attn = sim.softmax(dim=-1)
+        del sim
 
         out = einsum('b i j, b j d -> b i d', attn, v)
+        del attn, v
         out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+        del h
         return self.to_out(out)