slight efficiency gain by using += in attention.py

cognidesign · Sep 11, 2022 · 7708f4f · 7708f4f
1 parent b86a1de
commit 7708f4f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/ldm/modules/attention.py b/ldm/modules/attention.py
@@ -244,9 +244,9 @@ def forward(self, x, context=None):
 
     def _forward(self, x, context=None):
         x = x.contiguous() if x.device.type == 'mps' else x
-        x = self.attn1(self.norm1(x)) + x
-        x = self.attn2(self.norm2(x), context=context) + x
-        x = self.ff(self.norm3(x)) + x
+        x += self.attn1(self.norm1(x))
+        x += self.attn2(self.norm2(x), context=context)
+        x += self.ff(self.norm3(x))
         return x