From 625b5ff16dbfda240c5faf8f6154d87fca0d9dcb Mon Sep 17 00:00:00 2001
From: Hong Zhang <41229682+mi804@users.noreply.github.com>
Date: Tue, 24 Feb 2026 15:26:49 +0800
Subject: [PATCH] Apply suggestion from @gemini-code-assist[bot]

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 examples/qwen_image/model_training/train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/qwen_image/model_training/train.py b/examples/qwen_image/model_training/train.py
index 447c16a..3ebe524 100644
--- a/examples/qwen_image/model_training/train.py
+++ b/examples/qwen_image/model_training/train.py
@@ -156,7 +156,7 @@ if __name__ == "__main__":
         fp8_models=args.fp8_models,
         offload_models=args.offload_models,
         task=args.task,
-        device=accelerator.device,
+        device="cpu" if args.initialize_model_on_cpu else accelerator.device,
         zero_cond_t=args.zero_cond_t,
     )
     model_logger = ModelLogger(