diff --git a/opensora/models/diffusion/latte/modeling_latte.py b/opensora/models/diffusion/latte/modeling_latte.py
index 14d388275..f1b3a4f31 100644
--- a/opensora/models/diffusion/latte/modeling_latte.py
+++ b/opensora/models/diffusion/latte/modeling_latte.py
@@ -649,6 +649,7 @@ class LatteT2V(ModelMixin, ConfigMixin):
             attention_type: str = "default",
             caption_channels: int = None,
             video_length: int = 16,
+            attention_mode: str = 'flash'
     ):
         super().__init__()
         self.use_linear_projection = use_linear_projection
@@ -751,6 +752,7 @@ class LatteT2V(ModelMixin, ConfigMixin):
                     norm_elementwise_affine=norm_elementwise_affine,
                     norm_eps=norm_eps,
                     attention_type=attention_type,
+                    attention_mode=attention_mode
                 )
                 for d in range(num_layers)
             ]
@@ -775,6 +777,7 @@ class LatteT2V(ModelMixin, ConfigMixin):
                     norm_elementwise_affine=norm_elementwise_affine,
                     norm_eps=norm_eps,
                     attention_type=attention_type,
+                    attention_mode=attention_mode
                 )
                 for d in range(num_layers)
             ]
@@ -1161,4 +1164,4 @@ if __name__ == '__main__':
     model = LatteT2V.from_config(a)
     ckpt = torch.load(r"E:\下载\t2v.pt", map_location='cpu')['model']
     model.load_state_dict(ckpt)
-    print(model)
\ No newline at end of file
+    print(model)