Merge pull request #84 from beringresearch/tf-2.4_compat

Tf 2.4 compat
beringresearch · Jan 5, 2021 · 7712105 · 7712105
2 parents 0d21ad9 + 270f7c1
commit 7712105
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 50 deletions.
diff --git a/ivis/ivis.py b/ivis/ivis.py
@@ -2,7 +2,7 @@
 import json
 import os
 import shutil
-import pickle as pkl
+import dill as pkl
 import tensorflow as tf
 import numpy as np
 
@@ -406,7 +406,7 @@ def save_model(self, folder_path, overwrite=False):
 
         # save optimizer state in numpy array
         np.save(os.path.join(folder_path, 'optimizer_state.npy'),
-                self.model_.optimizer.get_weights())
+                np.array(self.model_.optimizer.get_weights(), dtype=object))
 
         json.dump(self.__getstate__(),
                   open(os.path.join(folder_path, 'ivis_params.json'), 'w'))

diff --git a/requirements.txt b/requirements.txt
@@ -3,4 +3,4 @@ numpy
 scikit-learn>0.20.0
 tqdm
 annoy>=1.15.2
-pytest
+dill
diff --git a/setup.py b/setup.py
@@ -23,7 +23,8 @@
           'numpy',
           'scikit-learn>0.20.0',
           'annoy>=1.15.2',
-          'tqdm'
+          'tqdm',
+          'dill'
       ],
       extras_require={
           'tests': ['pytest'],

diff --git a/tests/integration/test_h5.py b/tests/integration/test_h5.py
-Original file line number
+Diff line change
@@ Expand Up / @@ -3,4 +3,4 @@ numpy @@
     scikit-learn>0.20.0
     tqdm
     annoy>=1.15.2
-    pytest
+    dill