diff --git a/requirements.txt b/requirements.txt
index b904be5..a9d26dd 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -8,5 +8,4 @@ matplotlib
 flake8
 pytest
 black
-tensorflow>=2.1.0
 gpflow>=2.0.0
diff --git a/setup.py b/setup.py
index 880e884..4050d05 100644
--- a/setup.py
+++ b/setup.py
@@ -18,7 +18,7 @@
 
 setuptools.setup(
     name="pygpso",
-    version="0.5",
+    version="0.5.1",
     description="Bayesian optimisation method leveraging Gaussian Processes "
     "surrogate",
     long_description=long_description,
diff --git a/tests/test_gp_surrogate.py b/tests/test_gp_surrogate.py
index 8d3e481..57e268e 100644
--- a/tests/test_gp_surrogate.py
+++ b/tests/test_gp_surrogate.py
@@ -321,10 +321,12 @@ def test_save_load(self):
         loaded = GPRSurrogate.from_saved(self.TEMP_FOLDER)
         self.assertTrue(isinstance(loaded, GPRSurrogate))
         # assert GPR parameters
-        self.assertDictEqual(
-            gpflow.utilities.parameter_dict(self.gp_surr.gpflow_model),
-            gpflow.utilities.parameter_dict(loaded.gpflow_model),
-        )
+        for (key_orig, val_orig), (key_loaded, val_loaded) in zip(
+            gpflow.utilities.parameter_dict(self.gp_surr.gpflow_model).items(),
+            gpflow.utilities.parameter_dict(loaded.gpflow_model).items(),
+        ):
+            self.assertEqual(key_orig, key_loaded)
+            self.assertTrue((val_orig.numpy() == val_loaded.numpy()).all())
         # do test prediction with GPR model
         test_points = [
             point