diff --git a/conda-env.yml b/conda-env.yml
index a247d98..ffce778 100644
--- a/conda-env.yml
+++ b/conda-env.yml
@@ -6,7 +6,7 @@ dependencies:
   - python=3.7.*
   - numpy=1.18.*
   - pandas=1.0.*
-  - pandas-profiling=2.4.*
+  - pandas-profiling=2.8.*
   - matplotlib=3.2.*
   - seaborn=0.10.*
   - scikit-learn=0.22.*
@@ -15,4 +15,4 @@ dependencies:
   - cudatoolkit=10.*
   - tqdm=4.43.*
   - ipython
-  - jupyter
\ No newline at end of file
+  - jupyter
diff --git a/notebooks/CBFV/cbfv/composition.py b/notebooks/CBFV/cbfv/composition.py
index 1ec971a..6396f0d 100644
--- a/notebooks/CBFV/cbfv/composition.py
+++ b/notebooks/CBFV/cbfv/composition.py
@@ -290,10 +290,10 @@ def generate_features(df, elem_prop='oliynyk',
 
     # get the column names
     cols = X.columns.values
-    # find the mean value of each column
+    # find the median value of each column
     median_values = X[cols].median()
-    # fill the missing values in each column with the columns mean value
-    X[cols] = X[cols].fillna(median_values.iloc[0])
+    # fill the missing values in each column with the column's median value
+    X[cols] = X[cols].fillna(median_values)
     return X, y, formulae, skipped