From 096fd2eab312aa0a81522622bd8803d69afa5d5f Mon Sep 17 00:00:00 2001
From: "luyulei.233" <luyulei.233@bytedance.com>
Date: Sun, 29 Sep 2024 21:20:49 +0800
Subject: [PATCH] =?UTF-8?q?=E5=BC=80=E5=90=AFdeepspeed=E6=97=B6bfloat16?=
 =?UTF-8?q?=E7=B1=BB=E5=9E=8B=E9=94=99=E8=AF=AF?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 finetune/train_cogvideox_lora.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/finetune/train_cogvideox_lora.py b/finetune/train_cogvideox_lora.py
index 137f322..5ad4750 100644
--- a/finetune/train_cogvideox_lora.py
+++ b/finetune/train_cogvideox_lora.py
@@ -1042,7 +1042,7 @@ def main(args):
             "bf16" in accelerator.state.deepspeed_plugin.deepspeed_config
             and accelerator.state.deepspeed_plugin.deepspeed_config["bf16"]["enabled"]
         ):
-            weight_dtype = torch.float16
+            weight_dtype = torch.bfloat16
     else:
         if accelerator.mixed_precision == "fp16":
             weight_dtype = torch.float16