Fixed warning and added safe globals

torch_frame/__init__.py

@@ -12,13 +12,27 @@
     embedding,
 )
 from .data import TensorFrame
-from .typing import TaskType, Metric, DataFrame, NAStrategy
+from .typing import (
+    TaskType,
+    Metric,
+    DataFrame,
+    NAStrategy,
+    WITH_PT24,
+)
 from torch_frame.utils import save, load, cat  # noqa
 import torch_frame.data  # noqa
 import torch_frame.datasets  # noqa
 import torch_frame.nn  # noqa
 import torch_frame.gbdt  # noqa
 
+if WITH_PT24:
+    import torch
+
+    torch.serialization.add_safe_globals([
+        stype,
+        torch_frame.data.stats.StatType,
+    ])
+
 __version__ = '0.2.3'
 
 __all__ = [

torch_frame/typing.py

@@ -4,11 +4,15 @@
 from typing import Dict, List, Mapping, Union
 
 import pandas as pd
+import torch
 from torch import Tensor
 
 from torch_frame.data.multi_embedding_tensor import MultiEmbeddingTensor
 from torch_frame.data.multi_nested_tensor import MultiNestedTensor
 
+WITH_PT20 = int(torch.__version__.split('.')[0]) >= 2
+WITH_PT24 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 4
+
 
 class Metric(Enum):
     r"""The metric.

torch_frame/utils/io.py

@@ -80,7 +80,8 @@ def save(tensor_frame: TensorFrame,
 
 
 def load(
-    path: str, device: torch.device | None = None
+    path: str,
+    device: torch.device | None = None,
 ) -> tuple[TensorFrame, dict[str, dict[StatType, Any]] | None]:
     r"""Load saved :class:`TensorFrame` object and optional :obj:`col_stats`
     from a specified path.
@@ -95,7 +96,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)