[Bugfix] Place empty Qwen encoder-DP embeddings on the communication device (#31663)

2026-07-21 19:54:36 +08:00
parent 7c9257529f
commit 303896a475
1 changed files with 2 additions and 2 deletions
@@ -707,7 +707,7 @@ def run_dp_sharded_mrope_vision_model(
            out_dim = getattr(vision_model.config, "hidden_size", None)
            image_embeds_local = torch.empty(
                (0, embed_dim_reduction_factor, out_dim),
-                device=input_device,
+                device=vision_model.device,
                dtype=input_dtype,
            )
    else:
@@ -725,7 +725,7 @@ def run_dp_sharded_mrope_vision_model(
                out_dim = vision_model.config.hidden_size
            image_embeds_local = torch.empty(
                (0, out_dim),
-                device=input_device,
+                device=vision_model.device,
                dtype=input_dtype,
            )