Unverified Commit ad977635 authored by Harry Mellor's avatar Harry Mellor Committed by GitHub
Browse files

Set `torch_dtype` in `TransformersModel` (#13088)


Signed-off-by: default avatarHarry Mellor <19981378+hmellor@users.noreply.github.com>
parent 75e6e145
......@@ -143,6 +143,7 @@ class TransformersModel(nn.Module):
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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment