[CI]revert initialize_model context manager (#38426)

Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>
Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>
This commit is contained in:
Kunshang Ji
2026-03-29 00:56:50 +08:00
committed by GitHub
parent af89140efc
commit aa4eb0db78

View File

@@ -50,10 +50,14 @@ class BaseModelLoader(ABC):
device_config.device if load_config.device is None else load_config.device
)
target_device = torch.device(load_device)
with set_default_torch_dtype(model_config.dtype), target_device:
model = initialize_model(
vllm_config=vllm_config, model_config=model_config, prefix=prefix
)
with set_default_torch_dtype(model_config.dtype):
with target_device:
model = initialize_model(
vllm_config=vllm_config,
model_config=model_config,
prefix=prefix,
)
log_model_inspection(model)
logger.debug("Loading weights on %s ...", load_device)