diff --git a/paltax/input_pipeline.py b/paltax/input_pipeline.py
index 415a10a..1d07fb3 100644
--- a/paltax/input_pipeline.py
+++ b/paltax/input_pipeline.py
@@ -531,7 +531,7 @@ def initialize_cosmology_params(
     r_max = cosmology_utils.lagrangian_radius(cosmology_params,
                                               m_max * 10)
     cosmology_params = cosmology_utils.add_lookup_tables_to_cosmology_params(
-        cosmology_params_init, 1.5, dz / 2, r_min, r_max, 10000)
+        cosmology_params_init, max_source_z, dz / 2, r_min, r_max, 10000)
     extrenal_los_params = {'m_min': m_min, 'm_max': m_max, 'dz': dz}
     cosmology_params = los.add_los_lookup_tables_to_cosmology_params(
         extrenal_los_params, cosmology_params, max_source_z
diff --git a/paltax/input_pipeline_test.py b/paltax/input_pipeline_test.py
index eba2ca4..3d2c320 100644
--- a/paltax/input_pipeline_test.py
+++ b/paltax/input_pipeline_test.py
@@ -630,6 +630,7 @@ def test_initialize_cosmology_params(self):
         r_max = cosmology_utils.lagrangian_radius(cosmology_params, 1e11)
         # Cosmological values are calculated in incriments of dz / 2.
         self.assertAlmostEqual(cosmology_params['dz'], dz / 2)
+        self.assertAlmostEqual(cosmology_params['z_lookup_max'], 2.0)
         self.assertAlmostEqual(cosmology_params['r_min'], r_min, places=6)
         self.assertAlmostEqual(cosmology_params['r_max'], r_max, places=6)