diff --git a/xllm/core/runtime/dit_forward_params.h b/xllm/core/runtime/dit_forward_params.h
index d5e8e72d75..bc926bcce5 100644
--- a/xllm/core/runtime/dit_forward_params.h
+++ b/xllm/core/runtime/dit_forward_params.h
@@ -205,28 +205,30 @@ struct DiTForwardInput {
     }
 
     if (images.defined()) {
-      input.images = images.to(device, /*dtype=*/torch::kUInt8);
+      input.images = images.to(torch::kCPU, /*dtype=*/torch::kUInt8);
     }
 
     if (mask_images.defined()) {
-      input.mask_images = mask_images.to(device, /*dtype=*/torch::kUInt8);
+      input.mask_images = mask_images.to(torch::kCPU, /*dtype=*/torch::kUInt8);
     }
 
     for (auto& img : input.images_list) {
-      img = img.to(device, /*dtype=*/torch::kUInt8);
+      img = img.to(torch::kCPU, /*dtype=*/torch::kUInt8);
     }
 
     if (control_image.defined()) {
-      input.control_image = control_image.to(device, /*dtype=*/torch::kUInt8);
+      input.control_image =
+          control_image.to(torch::kCPU, /*dtype=*/torch::kUInt8);
     }
 
     if (last_images.defined()) {
-      input.last_images = last_images.to(device, /*dtype=*/torch::kUInt8);
+      input.last_images = last_images.to(torch::kCPU, /*dtype=*/torch::kUInt8);
     }
 
     if (prompt_audio.defined()) {
       input.prompt_audio = prompt_audio.to(device, torch::kFloat32);
     }
+
     return input;
   }