diff --git a/.github/imgs/multi_point_mut_performance_PGym_DCA.png b/.github/imgs/multi_point_mut_performance_PGym_DCA.png
index bd11b06..bc26a3a 100644
Binary files a/.github/imgs/multi_point_mut_performance_PGym_DCA.png and b/.github/imgs/multi_point_mut_performance_PGym_DCA.png differ
diff --git a/.github/imgs/single_point_mut_performance_PGyM_DCA.png b/.github/imgs/single_point_mut_performance_PGyM_DCA.png
index a7804a9..7009785 100644
Binary files a/.github/imgs/single_point_mut_performance_PGyM_DCA.png and b/.github/imgs/single_point_mut_performance_PGyM_DCA.png differ
diff --git a/scripts/ProteinGym_runs/run_performance_tests_proteingym_data.py b/scripts/ProteinGym_runs/run_performance_tests_proteingym_data.py
index 9d3a2bd..3a1d033 100644
--- a/scripts/ProteinGym_runs/run_performance_tests_proteingym_data.py
+++ b/scripts/ProteinGym_runs/run_performance_tests_proteingym_data.py
@@ -111,7 +111,7 @@ def plot_performance(mut_data, plot_name, mut_sep=':'):
             f'0' + r'$\rightarrow$' + f'{len(fitnesses)}', 
             color='tab:blue', size=4, ha='right'
         ))
-        assert len(x_dca) == len(fitnesses) == len(variants) == len(sequences)
+        assert len(x_dca) == len(fitnesses) == len(variants) == len(sequences) == len(y_pred)
         hybrid_perfs = []
         for i_t, train_size in enumerate([25, 50, 75, 100, 200]):
             try: