diff --git a/PWGHF/Core/HfMlResponseDsToKKPi.h b/PWGHF/Core/HfMlResponseDsToKKPi.h
index d63340a3752..1b27d52adde 100644
--- a/PWGHF/Core/HfMlResponseDsToKKPi.h
+++ b/PWGHF/Core/HfMlResponseDsToKKPi.h
@@ -231,6 +231,9 @@ class HfMlResponseDsToKKPi : public HfMlResponse<TypeOutputScore>
       FILL_MAP_DS(impactParameterXY0),
       FILL_MAP_DS(impactParameterXY1),
       FILL_MAP_DS(impactParameterXY2),
+      FILL_MAP_DS(impactParameterZ0),
+      FILL_MAP_DS(impactParameterZ1),
+      FILL_MAP_DS(impactParameterZ2),
       // TPC PID variables
       FILL_MAP_DS(nSigTpcPi0),
       FILL_MAP_DS(nSigTpcPi1),