Skip to content
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.

Sign up for GitHub

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

Fixed warning and added safe globals #423

Merged
merged 7 commits into from
Aug 24, 2024
+25 −4
Merged

Fixed warning and added safe globals #423

Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
fbcba9d
Issue 422: Fixed warning and added safe globals
NeelKondapalli Jul 31, 2024
9ef2da0
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 31, 2024
4ce8e5c
Removed serialization.add_safe_globals logic not available in torch 2…
NeelKondapalli Aug 1, 2024
d9c2b1a
Merge branch 'master' into load-warning-comply
akihironitta Aug 24, 2024
f072bcf
update
akihironitta Aug 24, 2024
54dab53
update
akihironitta Aug 24, 2024
febaeae
update
akihironitta Aug 24, 2024
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
5 changes: 4 additions & 1 deletion torch_frame/utils/io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,9 @@
from torch_frame.data.stats import StatType
from torch_frame.typing import TensorData

torch.serialization.add_safe_globals([torch_frame.stype])
torch.serialization.add_safe_globals([StatType])


def serialize_feat_dict(
feat_dict: dict[torch_frame.stype, TensorData]
Expand Down Expand Up @@ -95,7 +98,7 @@ def load(
tuple: A tuple of loaded :class:`TensorFrame` object and
optional :obj:`col_stats`.
"""
tf_dict, col_stats = torch.load(path)
tf_dict, col_stats = torch.load(path, weights_only=True)
tf_dict['feat_dict'] = deserialize_feat_dict(
tf_dict.pop('feat_serialized_dict'))
tensor_frame = TensorFrame(**tf_dict)
Expand Down
Loading