diff --git a/examples/qwen_image/model_training/validate_full/Qwen-Image-Edit.py b/examples/qwen_image/model_training/validate_full/Qwen-Image-Edit.py
index 4d56ee8..c08b485 100644
--- a/examples/qwen_image/model_training/validate_full/Qwen-Image-Edit.py
+++ b/examples/qwen_image/model_training/validate_full/Qwen-Image-Edit.py
@@ -15,6 +15,7 @@ pipe = QwenImagePipeline.from_pretrained(
     processor_config=ModelConfig(model_id="Qwen/Qwen-Image-Edit", origin_file_pattern="processor/"),
 )
 state_dict = load_state_dict("models/train/Qwen-Image-Edit_full/epoch-1.safetensors")
+pipe.dit.load_state_dict(state_dict)
 
 prompt = "将裙子改为粉色"
 image = Image.open("data/example_image_dataset/edit/image1.jpg").resize((1024, 1024))