diff --git a/nn_regression.ipynb b/nn_regression.ipynb
index f0eb9bb7e713348f1682dc8f1136a35529feffad..fc300b1b031ac9570d93bb9824b513d61d491730 100644
--- a/nn_regression.ipynb
+++ b/nn_regression.ipynb
@@ -72,11 +72,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2021-01-29T16:59:18.160658Z",
-     "start_time": "2021-01-29T16:59:18.149111Z"
+     "end_time": "2021-01-29T17:13:47.825250Z",
+     "start_time": "2021-01-29T17:13:47.815813Z"
     },
     "scrolled": true
    },
@@ -87,8 +87,6 @@
     "\n",
     "import warnings\n",
     "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
-    "warnings.filterwarnings(\"ignore\", message=\"numpy.dtype size changed\")\n",
-    "warnings.filterwarnings(\"ignore\", message=\"numpy.ufunc size changed\")\n",
     "\n",
     "import matplotlib.pyplot as plt\n",
     "import seaborn as sns\n",
@@ -107,9 +105,6 @@
     "import numpy as np\n",
     "from collections import Counter\n",
     "\n",
-    "# pymatgen\n",
-    "from pymatgen.core.periodic_table import Element\n",
-    "\n",
     "# sklearn\n",
     "from sklearn.metrics import r2_score, mean_absolute_error\n",
     "from sklearn.model_selection import KFold\n",
@@ -120,7 +115,10 @@
     "import pandas as pd\n",
     "\n",
     "# json\n",
-    "import json"
+    "import json\n",
+    "\n",
+    "# mendeleev\n",
+    "from mendeleev import element as Element"
    ]
   },
   {
@@ -441,7 +439,7 @@
     "elements_of_interest_unique = np.array(list(set(elements_of_interest)), dtype=object)\n",
     "Z_values = []\n",
     "for element in elements_of_interest_unique:\n",
-    "    Z_values.append(Element(element).Z)\n",
+    "    Z_values.append(Element(element).atomic_number)\n",
     "sorted_indices = np.argsort(Z_values)\n",
     "elements_of_interest_sorted = elements_of_interest_unique[sorted_indices]\n",
     "print(\"\\nThe following elements (in total {}) appear in the dataset:\\n\\n {}\".format(len(elements_of_interest_sorted), elements_of_interest_sorted))\n",
diff --git a/setup.py b/setup.py
index 33e5bc18ff2cee2afeb4c661949173ebe462346a..761255956f7b72fb3d4e1ee74f95e62b61616793 100644
--- a/setup.py
+++ b/setup.py
@@ -16,5 +16,5 @@ setup(
     description=metainfo['title'],
     long_description=metainfo['description'],
     packages=find_packages(),
-    install_requires=['tensorflow', 'keras', 'numpy', 'scipy', 'matplotlib', 'pandas', 'seaborn', 'pymatgen', 'sklearn'],
+    install_requires=['mendeleev', 'tensorflow', 'keras', 'numpy', 'scipy', 'matplotlib', 'pandas', 'seaborn', 'pymatgen', 'sklearn'],
 )