From ed0a9dd9b06817c9488b3147552f6b227bdafeb8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=AE=A3=E6=BA=90?= <xuanyuan.lb@alibaba-inc.com>
Date: Thu, 12 Sep 2024 15:51:49 +0800
Subject: [PATCH] add swa

---
 cogvideox/models/transformer3d.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/cogvideox/models/transformer3d.py b/cogvideox/models/transformer3d.py
index bf81079..21c1405 100644
--- a/cogvideox/models/transformer3d.py
+++ b/cogvideox/models/transformer3d.py
@@ -488,9 +488,6 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin):
                     height // p, 
                     width // p,
                     **ckpt_kwargs,
-                    num_frames=num_frames, 
-                    height=height // p, 
-                    width=width // p,
                 )
             else:
                 hidden_states, encoder_hidden_states = block(
@@ -498,6 +495,9 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin):
                     encoder_hidden_states=encoder_hidden_states,
                     temb=emb,
                     image_rotary_emb=image_rotary_emb,
+                    num_frames=num_frames, 
+                    height=height // p, 
+                    width=width // p,
                 )
 
         if not self.config.use_rotary_positional_embeddings: