remove dill dependency for issue in python 3.8 with dill>=0.3.6, use …

…partials for pickling. Add torch version cap for segfault issue

contextualized/tests.py

@@ -90,13 +90,17 @@ def test_save_load(self):
         C = np.random.uniform(0, 1, size=(100, 2))
         X = np.random.uniform(0, 1, size=(100, 2))
         Y = np.random.uniform(0, 1, size=(100, 2))
+        C2 = np.random.uniform(0, 1, size=(100, 2))
+        X2 = np.random.uniform(0, 1, size=(100, 2))
+        Y2 = np.random.uniform(0, 1, size=(100, 2))
         mlp = MLP(2, 2, 50, 5)
         Y_pred = mlp(torch.Tensor(X)).detach().numpy()
         save(mlp, 'unittest_model.pt')
         del mlp
         mlp_loaded = load('unittest_model.pt')
         Y_pred_loaded = mlp_loaded(torch.Tensor(X)).detach().numpy()
         assert np.all(Y_pred == Y_pred_loaded)
+        os.remove('unittest_model.pt')
 
         model = ContextualizedRegressor()
         model.fit(C, X, Y)
@@ -107,6 +111,15 @@ def test_save_load(self):
         Y_pred_loaded = model_loaded.predict(C, X)
         assert np.all(Y_pred == Y_pred_loaded)
         os.remove('unittest_model.pt')
+        model_loaded.fit(C2, X2, Y2)
+        Y_pred2 = model_loaded.predict(C2, X2)
+        assert not np.all(Y_pred_loaded == Y_pred2)
+        save(model_loaded, 'unittest_model.pt')
+        del model_loaded
+        model_loaded2 = load('unittest_model.pt')
+        Y_pred_loaded2 = model_loaded2.predict(C2, X2)
+        assert np.all(Y_pred2 == Y_pred_loaded2)
+        os.remove('unittest_model.pt')
 
         model = ContextualizedBayesianNetworks()
         model.fit(C, X)
@@ -117,6 +130,15 @@ def test_save_load(self):
         pred_loaded = model_loaded.predict_networks(C)
         assert np.all(np.array(pred) == np.array(pred_loaded))
         os.remove('unittest_model.pt')
+        model_loaded.fit(C2, X2)
+        pred2 = model_loaded.predict_networks(C2)
+        assert not np.all(np.array(pred_loaded) == np.array(pred2))
+        save(model_loaded, 'unittest_model.pt')
+        del model_loaded
+        model_loaded2 = load('unittest_model.pt')
+        pred_loaded2 = model_loaded2.predict_networks(C2)
+        assert np.all(np.array(pred2) == np.array(pred_loaded2))
+        os.remove('unittest_model.pt')
 
         model = ContextualizedCorrelationNetworks()
         model.fit(C, X)
@@ -127,6 +149,15 @@ def test_save_load(self):
         pred_loaded = model_loaded.predict_correlation(C)
         assert np.all(np.array(pred) == np.array(pred_loaded))
         os.remove('unittest_model.pt')
+        model_loaded.fit(C2, X2)
+        pred2 = model_loaded.predict_correlation(C2)
+        assert not np.all(np.array(pred_loaded) == np.array(pred2))
+        save(model_loaded, 'unittest_model.pt')
+        del model_loaded
+        model_loaded2 = load('unittest_model.pt')
+        pred_loaded2 = model_loaded2.predict_correlation(C2)
+        assert np.all(np.array(pred2) == np.array(pred_loaded2))
+        os.remove('unittest_model.pt')
 
         model = BayesianNetwork()
         model.fit(X)
@@ -137,6 +168,15 @@ def test_save_load(self):
         pred_loaded = model_loaded.measure_mses(X)
         assert np.all(np.array(pred) == np.array(pred_loaded))
         os.remove('unittest_model.pt')
+        model_loaded.fit(X2)
+        pred2 = model_loaded.measure_mses(X2)
+        assert not np.all(np.array(pred_loaded) == np.array(pred2))
+        save(model_loaded, 'unittest_model.pt')
+        del model_loaded
+        model_loaded2 = load('unittest_model.pt')
+        pred_loaded2 = model_loaded2.measure_mses(X2)
+        assert np.all(np.array(pred2) == np.array(pred_loaded2))
+        os.remove('unittest_model.pt')
 
         model = CorrelationNetwork()
         model.fit(X)
@@ -147,6 +187,15 @@ def test_save_load(self):
         pred_loaded = model_loaded.measure_mses(X)
         assert np.all(np.array(pred) == np.array(pred_loaded))
         os.remove('unittest_model.pt')
+        model_loaded.fit(X2)
+        pred2 = model_loaded.measure_mses(X2)
+        assert not np.all(np.array(pred_loaded) == np.array(pred2))
+        save(model_loaded, 'unittest_model.pt')
+        del model_loaded
+        model_loaded2 = load('unittest_model.pt')
+        pred_loaded2 = model_loaded2.measure_mses(X2)
+        assert np.all(np.array(pred2) == np.array(pred_loaded2))
+        os.remove('unittest_model.pt')
 
 
 if __name__ == '__main__':

contextualized/utils.py

@@ -3,7 +3,6 @@
 """
 
 import torch
-import dill
 
 
 def save(model, path):
@@ -14,7 +13,7 @@ def save(model, path):
 
     """
     with open(path, "wb") as out_file:
-        torch.save(model, out_file, pickle_module=dill)
+        torch.save(model, out_file)
 
 
 def load(path):
@@ -24,7 +23,7 @@ def load(path):
 
     """
     with open(path, "rb") as in_file:
-        model = torch.load(in_file, pickle_module=dill)
+        model = torch.load(in_file)
     return model
 
 

pyproject.toml

@@ -26,13 +26,12 @@ keywords = [
 ]
 dependencies = [
     'lightning>=2.0.0',
-    'torch>=2.1.0',
+    'torch>=2.0.0,<2.2.0',
     'torchvision>=0.8.0',
     'numpy>=1.19.0',
     'pandas>=2.0.0',
     'scikit-learn>=1.0.0',
     'igraph>=0.11.0',
-    'dill>=0.3.3',
     'matplotlib>=3.3.0',
 ]