Skip to content

Commit e946260

Browse files
authored
use get_tensor in safe_open (#1696)
1 parent edb3055 commit e946260

File tree

1 file changed

+3
-8
lines changed

1 file changed

+3
-8
lines changed

vllm/model_executor/weight_utils.py

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -243,8 +243,8 @@ def hf_model_weights_iterator(
243243
for st_file in hf_weights_files:
244244
with safe_open(st_file, framework="pt") as f:
245245
for name in f.keys():
246-
param = f.get_slice(name)
247-
yield name, convert_pyslice_to_tensor(param)
246+
param = f.get_tensor(name)
247+
yield name, param
248248
else:
249249
for bin_file in hf_weights_files:
250250
state = torch.load(bin_file, map_location="cpu")
@@ -265,12 +265,7 @@ def convert_pyslice_to_tensor(x: Any) -> torch.Tensor:
265265
tensor first.
266266
"""
267267
if not isinstance(x, torch.Tensor):
268-
try:
269-
x = x[:]
270-
except IndexError:
271-
# IndexError happens when the tensor is empty.
272-
# transformer.h.0.attn.masked_bias is empty in some gpt2 models.
273-
return torch.Tensor()
268+
x = x[:]
274269
return x
275270

276271

0 commit comments

Comments
 (0)