diff --git a/src/egamma_tnp/utils/pileup.py b/src/egamma_tnp/utils/pileup.py
index 6896885b..9e92890b 100644
--- a/src/egamma_tnp/utils/pileup.py
+++ b/src/egamma_tnp/utils/pileup.py
@@ -51,7 +51,7 @@ def create_correction(pu_data_histogram, pu_mc_array, outfile=None, normalize_pu
 
     ceval = cset.to_evaluator()
     corr = ceval["Pileup"]
-    np.testing.assert_allclose(corr.evaluate(np.arange(100, dtype=float)), (pu_data / pu_mc), err_msg="Pileup correction does not match input data")
+    np.testing.assert_allclose(corr.evaluate(np.arange(len(pu_data), dtype=float)), (pu_data / pu_mc), err_msg="Pileup correction does not match input data")
     return correctionlib_wrapper(corr)