From 903daa8939dd7cfc3d57438f0cbb5af17cac21f0 Mon Sep 17 00:00:00 2001
From: Serge Panev <spanev@nvidia.com>
Date: Wed, 11 Sep 2024 18:32:43 +0200
Subject: [PATCH] Add more whitelisted variables for torch safe loading

---
 torch_frame/__init__.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/torch_frame/__init__.py b/torch_frame/__init__.py
index 7161acc38..3e35f4d68 100644
--- 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'