Set torch_dtype in TransformersModel (vllm-project#13088)

Signed-off-by: Harry Mellor <[email protected]>
yottalabsai · Feb 18, 2025 · 56b0f73 · 56b0f73
1 parent a0d55f4
commit 56b0f73
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/vllm/model_executor/models/transformers.py b/vllm/model_executor/models/transformers.py
@@ -143,6 +143,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         self.model: PreTrainedModel = AutoModel.from_config(
             self.config,
             attn_implementation="vllm",
+            torch_dtype=vllm_config.model_config.dtype,
             trust_remote_code=vllm_config.model_config.trust_remote_code,
         )
         prefix = self.model.base_model_prefix