Add more whitelisted variables for torch safe loading

pyg-team · Sep 11, 2024 · 903daa8 · 903daa8
1 parent 546f1a2
commit 903daa8
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/torch_frame/__init__.py b/torch_frame/__init__.py
@@ -27,10 +27,16 @@
 
 if WITH_PT24:
     import torch
+    import numpy as np
+    import _codecs
 
     torch.serialization.add_safe_globals([
         stype,
         torch_frame.data.stats.StatType,
+        np.core.multiarray.scalar,
+        np.dtype,
+        type(np.dtype(np.int32)),
+        _codecs.encode,
     ])
 
 __version__ = '0.2.3'