specified device_map=auto

3outeille · 3outeille · commit 2ac9219ba46d · 2025-11-03T12:32:09.000Z
diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py
@@ -234,7 +234,7 @@ def _test_model_dense_forward_impl(rank, mode):
 
     # Load non-TP model and move to same device as TP model
     device = model_tp.device
-    model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto")
+    model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto", device_map="auto")
     model = model.to(device)
 
     if mode == "eval":
