fix: 424 allow any np.intX as training target

dreamquark-ai · Apr 24, 2023 · 63a8dba · 63a8dba
1 parent fc59ea6
commit 63a8dba
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pytorch_tabnet/utils.py b/pytorch_tabnet/utils.py
@@ -422,8 +422,8 @@ def define_device(device_name):
 
 class ComplexEncoder(json.JSONEncoder):
     def default(self, obj):
-        if isinstance(obj, np.int64):
-            return int(obj)
+        if isinstance(obj, (np.generic, np.ndarray)):
+            return obj.tolist()
         # Let the base class default method raise the TypeError
         return json.JSONEncoder.default(self, obj)