Fix potential issue.

comfyanonymous · May 20, 2024 · 1900e51 · 1900e51
1 parent 276f8fc
commit 1900e51
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py
@@ -625,15 +625,15 @@ def forward(self, x, context=None, transformer_options={}):
         x = self.norm(x)
         if not self.use_linear:
             x = self.proj_in(x)
-        x = x.movedim(1, -1).flatten(1, 2).contiguous()
+        x = x.movedim(1, 3).flatten(1, 2).contiguous()
         if self.use_linear:
             x = self.proj_in(x)
         for i, block in enumerate(self.transformer_blocks):
             transformer_options["block_index"] = i
             x = block(x, context=context[i], transformer_options=transformer_options)
         if self.use_linear:
             x = self.proj_out(x)
-        x = x.reshape(x.shape[0], h, w, x.shape[-1]).movedim(-1, 1).contiguous()
+        x = x.reshape(x.shape[0], h, w, x.shape[-1]).movedim(3, 1).contiguous()
         if not self.use_linear:
             x = self.proj_out(x)
         return x + x_in