Skip to content

fix bytes2tensor error in torch version == 2.6.0. #950

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 1, 2025
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion lightllm/server/embed_cache/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ def tensor2bytes(t: torch.Tensor):

def bytes2tensor(b):
# return torch.from_numpy(np.frombuffer(b, dtype=np.float16)).cuda()
return torch.load(BytesIO(b))
return torch.load(BytesIO(b), weights_only=False)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Explicitly setting weights_only=False is the correct fix for the breaking change in upcoming PyTorch versions where this argument will default to True. However, this setting can introduce a security vulnerability if the serialized data comes from an untrusted source, as it can lead to arbitrary code execution during deserialization.

In this specific case, the data is generated internally by the tensor2bytes function and passed via shared memory, so it can be considered trusted. To improve code clarity and prevent future misuse by developers who may not be aware of the context, I recommend adding a comment explaining why weights_only=False is necessary and why it's considered safe here.

Suggested change
return torch.load(BytesIO(b), weights_only=False)
return torch.load(BytesIO(b), weights_only=False) # Required for unpickling tensor objects; safe as data is generated internally.



def create_shm(name, data):
Expand Down