diff --git a/Tools/PIDML/KaonPidTask.cxx b/Tools/PIDML/KaonPidTask.cxx
index acc29e47bee..3105292d44c 100644
--- a/Tools/PIDML/KaonPidTask.cxx
+++ b/Tools/PIDML/KaonPidTask.cxx
@@ -46,7 +46,7 @@ struct KaonPidTask {
   SliceCache cache;
   Preslice<aod::Tracks> perCol = aod::track::collisionId;
 
-  std::shared_ptr<PidONNXModel> pidModel; // creates a shared pointer to a new instance 'pidmodel'.
+  std::shared_ptr<PidONNXModel<o2::aod::MyTracks>> pidModel; // creates a shared pointer to a new instance 'pidmodel'.
   HistogramRegistry histos{"Histos", {}, OutputObjHandlingPolicy::AnalysisObject};
 
   Configurable<float> cfgZvtxCut{"cfgZvtxCut", 10, "Z vtx cut"};
@@ -84,7 +84,7 @@ struct KaonPidTask {
     if (cfgUseCCDB) {
       ccdbApi.init(cfgCCDBURL); // Initializes ccdbApi when cfgUseCCDB is set to 'true'
     }
-    pidModel = std::make_shared<PidONNXModel>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, cfgTimestamp.value, cfgPid.value, cfgCertainty.value);
+    pidModel = std::make_shared<PidONNXModel<o2::aod::MyTracks>>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, cfgTimestamp.value, cfgPid.value, cfgCertainty.value);
 
     histos.add("hChargePos", ";z;", kTH1F, {{3, -1.5, 1.5}});
     histos.add("hChargeNeg", ";z;", kTH1F, {{3, -1.5, 1.5}});
diff --git a/Tools/PIDML/pidML.h b/Tools/PIDML/pidML.h
index 821da7efc42..6629b88475b 100644
--- a/Tools/PIDML/pidML.h
+++ b/Tools/PIDML/pidML.h
@@ -19,6 +19,8 @@
 
 #include "Framework/AnalysisDataModel.h"
 #include "Common/DataModel/PIDResponse.h"
+#include "Common/DataModel/Centrality.h"
+#include "Common/DataModel/Multiplicity.h"
 
 namespace o2::aod
 {
diff --git a/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx b/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx
index fe4d591e590..8afede63e86 100644
--- a/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx
+++ b/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx
@@ -66,7 +66,6 @@ struct PidMlBatchEffAndPurProducer {
   std::array<std::shared_ptr<TH1>, kNPids> hMCPositive;
 
   o2::ccdb::CcdbApi ccdbApi;
-  std::vector<PidONNXModel> models;
 
   Configurable<std::vector<int32_t>> cfgPids{"pids", std::vector<int32_t>(kPids, kPids + kNPids), "PIDs to predict"};
   Configurable<std::array<double, kNDetectors>> cfgDetectorsPLimits{"detectors-p-limits", std::array<double, kNDetectors>(pidml_pt_cuts::defaultModelPLimits), "\"use {detector} when p >= y_{detector}\": array of 3 doubles [y_TPC, y_TOF, y_TRD]"};
@@ -82,6 +81,7 @@ struct PidMlBatchEffAndPurProducer {
   using BigTracks = soa::Filtered<soa::Join<aod::FullTracks, aod::TracksDCA, aod::pidTOFbeta, aod::TrackSelection, aod::TOFSignal, aod::McTrackLabels,
                                             aod::pidTPCFullPi, aod::pidTPCFullKa, aod::pidTPCFullPr, aod::pidTPCFullEl, aod::pidTPCFullMu,
                                             aod::pidTOFFullPi, aod::pidTOFFullKa, aod::pidTOFFullPr, aod::pidTOFFullEl, aod::pidTOFFullMu>>;
+  std::vector<PidONNXModel<BigTracks>> models;
 
   void initHistos()
   {
@@ -203,11 +203,11 @@ struct PidMlBatchEffAndPurProducer {
     if (cfgUseCCDB && bc.runNumber() != currentRunNumber) {
       uint64_t timestamp = cfgUseFixedTimestamp ? cfgTimestamp.value : bc.timestamp();
       for (const int32_t& pid : cfgPids.value)
-        models.emplace_back(PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value,
+        models.emplace_back(PidONNXModel<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value,
                                          ccdbApi, timestamp, pid, 1.1, &cfgDetectorsPLimits.value[0]));
     } else {
       for (int32_t& pid : cfgPids.value)
-        models.emplace_back(PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value,
+        models.emplace_back(PidONNXModel<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value,
                                          ccdbApi, -1, pid, 1.1, &cfgDetectorsPLimits.value[0]));
     }
 
diff --git a/Tools/PIDML/pidMLEffAndPurProducer.cxx b/Tools/PIDML/pidMLEffAndPurProducer.cxx
index 7e626a8d631..620fa1b56b1 100644
--- a/Tools/PIDML/pidMLEffAndPurProducer.cxx
+++ b/Tools/PIDML/pidMLEffAndPurProducer.cxx
@@ -24,7 +24,6 @@
 #include "Common/DataModel/TrackSelectionTables.h"
 #include "Common/DataModel/PIDResponse.h"
 #include "Tools/PIDML/pidOnnxModel.h"
-#include "pidOnnxModel.h"
 #include "Tools/PIDML/pidUtils.h"
 
 using namespace o2;
@@ -35,7 +34,6 @@ using namespace pidml::pidutils;
 struct PidMlEffAndPurProducer {
   HistogramRegistry histos{"histos", {}, OutputObjHandlingPolicy::AnalysisObject};
 
-  PidONNXModel pidModel;
   Configurable<int> cfgPid{"pid", 211, "PID to predict"};
   Configurable<double> cfgNSigmaCut{"n-sigma-cut", 3.0f, "TPC and TOF PID nSigma cut"};
   Configurable<std::array<double, kNDetectors>> cfgDetectorsPLimits{"detectors-p-limits", std::array<double, kNDetectors>(pidml_pt_cuts::defaultModelPLimits), "\"use {detector} when p >= y_{detector}\": array of 3 doubles [y_TPC, y_TOF, y_TRD]"};
@@ -57,6 +55,7 @@ struct PidMlEffAndPurProducer {
   using BigTracks = soa::Filtered<soa::Join<aod::FullTracks, aod::TracksDCA, aod::pidTOFbeta, aod::TrackSelection, aod::TOFSignal, aod::McTrackLabels,
                                             aod::pidTPCFullPi, aod::pidTPCFullKa, aod::pidTPCFullPr, aod::pidTPCFullEl, aod::pidTPCFullMu,
                                             aod::pidTOFFullPi, aod::pidTOFFullKa, aod::pidTOFFullPr, aod::pidTOFFullEl, aod::pidTOFFullMu>>;
+  PidONNXModel<BigTracks> pidModel;
 
   typedef struct nSigma_t {
     double tpc, tof;
@@ -116,7 +115,7 @@ struct PidMlEffAndPurProducer {
     if (cfgUseCCDB) {
       ccdbApi.init(cfgCCDBURL);
     } else {
-      pidModel = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1,
+      pidModel = PidONNXModel<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1,
                               cfgPid.value, cfgCertainty.value, &cfgDetectorsPLimits.value[0]);
     }
 
@@ -153,7 +152,7 @@ struct PidMlEffAndPurProducer {
     auto bc = collisions.iteratorAt(0).bc_as<aod::BCsWithTimestamps>();
     if (cfgUseCCDB && bc.runNumber() != currentRunNumber) {
       uint64_t timestamp = cfgUseFixedTimestamp ? cfgTimestamp.value : bc.timestamp();
-      pidModel = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp,
+      pidModel = PidONNXModel<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp,
                               cfgPid.value, cfgCertainty.value, &cfgDetectorsPLimits.value[0]);
     }
 
diff --git a/Tools/PIDML/pidOnnxInterface.h b/Tools/PIDML/pidOnnxInterface.h
index 03e7bc19146..9aa0a5c6fcb 100644
--- a/Tools/PIDML/pidOnnxInterface.h
+++ b/Tools/PIDML/pidOnnxInterface.h
@@ -45,6 +45,7 @@ static const std::vector<std::string> cutVarLabels = {
 
 } // namespace pidml_pt_cuts
 
+template <typename T>
 struct PidONNXInterface {
   PidONNXInterface(std::string& localPath, std::string& ccdbPath, bool useCCDB, o2::ccdb::CcdbApi& ccdbApi, uint64_t timestamp, std::vector<int> const& pids, o2::framework::LabeledArray<double> const& pLimits, std::vector<double> const& minCertainties, bool autoMode) : mNPids{pids.size()}, mPLimits{pLimits}
   {
@@ -78,8 +79,7 @@ struct PidONNXInterface {
   PidONNXInterface& operator=(const PidONNXInterface&) = delete;
   ~PidONNXInterface() = default;
 
-  template <typename T>
-  float applyModel(const T& track, int pid)
+  float applyModel(const T::iterator& track, int pid)
   {
     for (std::size_t i = 0; i < mNPids; i++) {
       if (mModels[i].mPid == pid) {
@@ -90,8 +90,7 @@ struct PidONNXInterface {
     return -1.0f;
   }
 
-  template <typename T>
-  bool applyModelBoolean(const T& track, int pid)
+  bool applyModelBoolean(const T::iterator& track, int pid)
   {
     for (std::size_t i = 0; i < mNPids; i++) {
       if (mModels[i].mPid == pid) {
@@ -110,7 +109,7 @@ struct PidONNXInterface {
     minCertainties = std::vector<double>(mNPids, 0.5);
   }
 
-  std::vector<PidONNXModel> mModels;
+  std::vector<PidONNXModel<T>> mModels;
   std::size_t mNPids;
   o2::framework::LabeledArray<double> mPLimits;
 };
diff --git a/Tools/PIDML/pidOnnxModel.h b/Tools/PIDML/pidOnnxModel.h
index 207bc18cd8e..c1c9a59bfe0 100644
--- a/Tools/PIDML/pidOnnxModel.h
+++ b/Tools/PIDML/pidOnnxModel.h
@@ -17,12 +17,16 @@
 #ifndef TOOLS_PIDML_PIDONNXMODEL_H_
 #define TOOLS_PIDML_PIDONNXMODEL_H_
 
+#include <Framework/ASoA.h>
 #include <array>
+#include <algorithm>
 #include <cstdint>
+#include <cstring>
 #include <limits>
+#include <optional>
 #include <string>
-#include <algorithm>
 #include <map>
+#include <type_traits>
 #include <utility>
 #include <memory>
 #include <vector>
@@ -73,6 +77,7 @@ bool readJsonFile(const std::string& config, rapidjson::Document& d)
 }
 } // namespace
 
+template <typename T>
 struct PidONNXModel {
  public:
   PidONNXModel(std::string& localPath, std::string& ccdbPath, bool useCCDB, o2::ccdb::CcdbApi& ccdbApi, uint64_t timestamp,
@@ -135,14 +140,12 @@ struct PidONNXModel {
   PidONNXModel& operator=(const PidONNXModel&) = delete;
   ~PidONNXModel() = default;
 
-  template <typename T>
-  float applyModel(const T& track)
+  float applyModel(const typename T::iterator& track)
   {
     return getModelOutput(track);
   }
 
-  template <typename T>
-  bool applyModelBoolean(const T& track)
+  bool applyModelBoolean(const typename T::iterator& track)
   {
     return getModelOutput(track) >= mMinCertainty;
   }
@@ -203,7 +206,9 @@ struct PidONNXModel {
     LOG(info) << "Using configuration files: " << localTrainColumnsPath << ", " << localScalingParamsPath;
     if (readJsonFile(localTrainColumnsPath, trainColumnsDoc)) {
       for (auto& param : trainColumnsDoc["columns_for_training"].GetArray()) {
-        mTrainColumns.emplace_back(param.GetString());
+        auto columnLabel = param.GetString();
+        mTrainColumns.emplace_back(columnLabel);
+        mGetters.emplace_back(o2::soa::row_helpers::getColumnGetterByLabel<float, T>(columnLabel));
       }
     }
     if (readJsonFile(localScalingParamsPath, scalingParamsDoc)) {
@@ -213,52 +218,49 @@ struct PidONNXModel {
     }
   }
 
-  template <typename T>
-  std::vector<float> createInputsSingle(const T& track)
+  static float scale(float value, const std::pair<float, float>& scalingParams)
   {
-    // TODO: Hardcoded for now. Planning to implement RowView extension to get runtime access to selected columns
-    // sign is short, trackType and tpcNClsShared uint8_t
+    return (value - scalingParams.first) / scalingParams.second;
+  }
 
-    float scaledTPCSignal = (track.tpcSignal() - mScalingParams.at("fTPCSignal").first) / mScalingParams.at("fTPCSignal").second;
+  std::vector<float> getValues(const typename T::iterator& track)
+  {
+    std::vector<float> output;
+    output.reserve(mTrainColumns.size());
 
-    std::vector<float> inputValues{scaledTPCSignal};
+    bool useTOF = !tofMissing(track) && inPLimit(track, mPLimits[kTPCTOF]);
+    bool useTRD = !trdMissing(track) && inPLimit(track, mPLimits[kTPCTOFTRD]);
 
-    // When TRD Signal shouldn't be used we pass quiet_NaNs to the network
-    if (!inPLimit(track, mPLimits[kTPCTOFTRD]) || trdMissing(track)) {
-      inputValues.push_back(std::numeric_limits<float>::quiet_NaN());
-      inputValues.push_back(std::numeric_limits<float>::quiet_NaN());
-    } else {
-      float scaledTRDSignal = (track.trdSignal() - mScalingParams.at("fTRDSignal").first) / mScalingParams.at("fTRDSignal").second;
-      inputValues.push_back(scaledTRDSignal);
-      inputValues.push_back(track.trdPattern());
-    }
+    for (uint32_t i = 0; i < mTrainColumns.size(); ++i) {
+      auto& columnLabel = mTrainColumns[i];
 
-    // When TOF Signal shouldn't be used we pass quiet_NaNs to the network
-    if (!inPLimit(track, mPLimits[kTPCTOF]) || tofMissing(track)) {
-      inputValues.push_back(std::numeric_limits<float>::quiet_NaN());
-      inputValues.push_back(std::numeric_limits<float>::quiet_NaN());
-    } else {
-      float scaledTOFSignal = (track.tofSignal() - mScalingParams.at("fTOFSignal").first) / mScalingParams.at("fTOFSignal").second;
-      float scaledBeta = (track.beta() - mScalingParams.at("fBeta").first) / mScalingParams.at("fBeta").second;
-      inputValues.push_back(scaledTOFSignal);
-      inputValues.push_back(scaledBeta);
-    }
+      if (
+        ((columnLabel == "fTRDSignal" || columnLabel == "fTRDPattern") && !useTRD) ||
+        ((columnLabel == "fTOFSignal" || columnLabel == "fBeta") && !useTOF)) {
+        output.push_back(std::numeric_limits<float>::quiet_NaN());
+        continue;
+      }
+
+      std::optional<std::pair<float, float>> scalingParams = std::nullopt;
+
+      auto scalingParamsEntry = mScalingParams.find(columnLabel);
+      if (scalingParamsEntry != mScalingParams.end()) {
+        scalingParams = scalingParamsEntry->second;
+      }
 
-    float scaledX = (track.x() - mScalingParams.at("fX").first) / mScalingParams.at("fX").second;
-    float scaledY = (track.y() - mScalingParams.at("fY").first) / mScalingParams.at("fY").second;
-    float scaledZ = (track.z() - mScalingParams.at("fZ").first) / mScalingParams.at("fZ").second;
-    float scaledAlpha = (track.alpha() - mScalingParams.at("fAlpha").first) / mScalingParams.at("fAlpha").second;
-    float scaledTPCNClsShared = (static_cast<float>(track.tpcNClsShared()) - mScalingParams.at("fTPCNClsShared").first) / mScalingParams.at("fTPCNClsShared").second;
-    float scaledDcaXY = (track.dcaXY() - mScalingParams.at("fDcaXY").first) / mScalingParams.at("fDcaXY").second;
-    float scaledDcaZ = (track.dcaZ() - mScalingParams.at("fDcaZ").first) / mScalingParams.at("fDcaZ").second;
+      float value = mGetters[i](track);
 
-    inputValues.insert(inputValues.end(), {track.p(), track.pt(), track.px(), track.py(), track.pz(), static_cast<float>(track.sign()), scaledX, scaledY, scaledZ, scaledAlpha, static_cast<float>(track.trackType()), scaledTPCNClsShared, scaledDcaXY, scaledDcaZ});
+      if (scalingParams) {
+        value = scale(value, scalingParams.value());
+      }
+
+      output.push_back(value);
+    }
 
-    return inputValues;
+    return output;
   }
 
-  template <typename T>
-  float getModelOutput(const T& track)
+  float getModelOutput(const typename T::iterator& track)
   {
     // First rank of the expected model input is -1 which means that it is dynamic axis.
     // Axis is exported as dynamic to make it possible to run model inference with the batch of
@@ -268,7 +270,7 @@ struct PidONNXModel {
     auto input_shape = mInputShapes[0];
     input_shape[0] = batch_size;
 
-    std::vector<float> inputTensorValues = createInputsSingle(track);
+    std::vector<float> inputTensorValues = getValues(track);
     std::vector<Ort::Value> inputTensors;
 
 #if __has_include(<onnxruntime/core/session/onnxruntime_cxx_api.h>)
@@ -323,6 +325,7 @@ struct PidONNXModel {
   }
 
   std::vector<std::string> mTrainColumns;
+  std::vector<float(*)(const typename T::iterator&)> mGetters;
   std::map<std::string, std::pair<float, float>> mScalingParams;
 
   std::shared_ptr<Ort::Env> mEnv = nullptr;
diff --git a/Tools/PIDML/qaPidML.cxx b/Tools/PIDML/qaPidML.cxx
index 54cafc437ab..a34963e8f72 100644
--- a/Tools/PIDML/qaPidML.cxx
+++ b/Tools/PIDML/qaPidML.cxx
@@ -342,11 +342,6 @@ struct QaPidML {
     }
   }
 
-  // one model for one particle
-  PidONNXModel model211;
-  PidONNXModel model2212;
-  PidONNXModel model321;
-
   Configurable<std::string> cfgPathCCDB{"ccdb-path", "Users/m/mkabus/PIDML", "base path to the CCDB directory with ONNX models"};
   Configurable<std::string> cfgCCDBURL{"ccdb-url", "http://alice-ccdb.cern.ch", "URL of the CCDB repository"};
   Configurable<bool> cfgUseCCDB{"useCCDB", true, "Whether to autofetch ML model from CCDB. If false, local file will be used."};
@@ -355,26 +350,32 @@ struct QaPidML {
   o2::ccdb::CcdbApi ccdbApi;
   int currentRunNumber = -1;
 
+  Filter trackFilter = requireGlobalTrackInFilter();
+  using pidTracks = soa::Filtered<soa::Join<aod::Tracks, aod::TracksExtra, aod::McTrackLabels, aod::TracksDCA, aod::TrackSelection, aod::pidTOFbeta, aod::TOFSignal>>;
+
+  // one model for one particle
+  PidONNXModel<pidTracks> model211;
+  PidONNXModel<pidTracks> model2212;
+  PidONNXModel<pidTracks> model321;
+
   void init(InitContext const&)
   {
     if (cfgUseCCDB) {
       ccdbApi.init(cfgCCDBURL);
     } else {
-      model211 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 211, 0.5f, pSwitchValue[0]);
-      model2212 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 2211, 0.5f, pSwitchValue[1]);
-      model321 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 321, 0.5f, pSwitchValue[2]);
+      model211 = PidONNXModel<pidTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 211, 0.5f, pSwitchValue[0]);
+      model2212 = PidONNXModel<pidTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 2211, 0.5f, pSwitchValue[1]);
+      model321 = PidONNXModel<pidTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 321, 0.5f, pSwitchValue[2]);
     }
   }
 
-  Filter trackFilter = requireGlobalTrackInFilter();
-  using pidTracks = soa::Filtered<soa::Join<aod::Tracks, aod::TracksExtra, aod::McTrackLabels, aod::TracksDCA, aod::TrackSelection, aod::pidTOFbeta, aod::TOFSignal>>;
   void process(aod::Collisions const& collisions, pidTracks const& tracks, aod::McParticles const& /*mcParticles*/, aod::BCsWithTimestamps const&)
   {
     auto bc = collisions.iteratorAt(0).bc_as<aod::BCsWithTimestamps>();
     if (cfgUseCCDB && bc.runNumber() != currentRunNumber) {
-      model211 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 211, 0.5f, pSwitchValue[0]);
-      model2212 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 2211, 0.5f, pSwitchValue[1]);
-      model321 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 321, 0.5f, pSwitchValue[2]);
+      model211 = PidONNXModel<pidTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 211, 0.5f, pSwitchValue[0]);
+      model2212 = PidONNXModel<pidTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 2211, 0.5f, pSwitchValue[1]);
+      model321 = PidONNXModel<pidTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 321, 0.5f, pSwitchValue[2]);
     }
 
     for (auto& track : tracks) {
diff --git a/Tools/PIDML/simpleApplyPidOnnxInterface.cxx b/Tools/PIDML/simpleApplyPidOnnxInterface.cxx
index 82f644103ed..d023f7390ee 100644
--- a/Tools/PIDML/simpleApplyPidOnnxInterface.cxx
+++ b/Tools/PIDML/simpleApplyPidOnnxInterface.cxx
@@ -39,8 +39,6 @@ DECLARE_SOA_TABLE(MlPidResults, "AOD", "MLPIDRESULTS", o2::soa::Index<>, mlpidre
 } // namespace o2::aod
 
 struct SimpleApplyOnnxInterface {
-  PidONNXInterface pidInterface; // One instance to manage all needed ONNX models
-
   Configurable<LabeledArray<double>> cfgPTCuts{"pT_cuts", {pidml_pt_cuts::cuts[0], pidml_pt_cuts::nPids, pidml_pt_cuts::nCutVars, pidml_pt_cuts::pidLabels, pidml_pt_cuts::cutVarLabels}, "pT cuts for each output pid and each detector configuration"};
   Configurable<std::vector<int>> cfgPids{"pids", std::vector<int>{pidml_pt_cuts::pids_v}, "PIDs to predict"};
   Configurable<std::vector<double>> cfgCertainties{"certainties", std::vector<double>{pidml_pt_cuts::certainties_v}, "Min certainties of the models to accept given particle to be of given kind"};
@@ -65,12 +63,14 @@ struct SimpleApplyOnnxInterface {
   // Filter on isGlobalTrack (TracksSelection)
   using BigTracks = soa::Filtered<soa::Join<aod::FullTracks, aod::TracksDCA, aod::pidTOFbeta, aod::TrackSelection, aod::TOFSignal>>;
 
+  PidONNXInterface<BigTracks> pidInterface; // One instance to manage all needed ONNX models
+
   void init(InitContext const&)
   {
     if (cfgUseCCDB) {
       ccdbApi.init(cfgCCDBURL);
     } else {
-      pidInterface = PidONNXInterface(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, cfgPids.value, cfgPTCuts.value, cfgCertainties.value, cfgAutoMode.value);
+      pidInterface = PidONNXInterface<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, cfgPids.value, cfgPTCuts.value, cfgCertainties.value, cfgAutoMode.value);
     }
   }
 
@@ -79,7 +79,7 @@ struct SimpleApplyOnnxInterface {
     auto bc = collisions.iteratorAt(0).bc_as<aod::BCsWithTimestamps>();
     if (cfgUseCCDB && bc.runNumber() != currentRunNumber) {
       uint64_t timestamp = cfgUseFixedTimestamp ? cfgTimestamp.value : bc.timestamp();
-      pidInterface = PidONNXInterface(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp, cfgPids.value, cfgPTCuts.value, cfgCertainties.value, cfgAutoMode.value);
+      pidInterface = PidONNXInterface<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp, cfgPids.value, cfgPTCuts.value, cfgCertainties.value, cfgAutoMode.value);
     }
 
     for (auto& track : tracks) {
diff --git a/Tools/PIDML/simpleApplyPidOnnxModel.cxx b/Tools/PIDML/simpleApplyPidOnnxModel.cxx
index e261c1e10bc..c2ba88e9b82 100644
--- a/Tools/PIDML/simpleApplyPidOnnxModel.cxx
+++ b/Tools/PIDML/simpleApplyPidOnnxModel.cxx
@@ -39,7 +39,6 @@ DECLARE_SOA_TABLE(MlPidResults, "AOD", "MLPIDRESULTS", o2::soa::Index<>, mlpidre
 } // namespace o2::aod
 
 struct SimpleApplyOnnxModel {
-  PidONNXModel pidModel; // One instance per model, e.g., one per each pid to predict
   Configurable<int> cfgPid{"pid", 211, "PID to predict"};
   Configurable<double> cfgCertainty{"certainty", 0.5, "Min certainty of the model to accept given particle to be of given kind"};
 
@@ -61,13 +60,14 @@ struct SimpleApplyOnnxModel {
   // TPC signal (FullTracks), TOF signal (TOFSignal), TOF beta (pidTOFbeta), dcaXY and dcaZ (TracksDCA)
   // Filter on isGlobalTrack (TracksSelection)
   using BigTracks = soa::Filtered<soa::Join<aod::FullTracks, aod::TracksDCA, aod::pidTOFbeta, aod::TrackSelection, aod::TOFSignal>>;
+  PidONNXModel<BigTracks> pidModel; // One instance per model, e.g., one per each pid to predict
 
   void init(InitContext const&)
   {
     if (cfgUseCCDB) {
       ccdbApi.init(cfgCCDBURL);
     } else {
-      pidModel = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, cfgPid.value, cfgCertainty.value);
+      pidModel = PidONNXModel<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, cfgPid.value, cfgCertainty.value);
     }
   }
 
@@ -76,7 +76,7 @@ struct SimpleApplyOnnxModel {
     auto bc = collisions.iteratorAt(0).bc_as<aod::BCsWithTimestamps>();
     if (cfgUseCCDB && bc.runNumber() != currentRunNumber) {
       uint64_t timestamp = cfgUseFixedTimestamp ? cfgTimestamp.value : bc.timestamp();
-      pidModel = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp, cfgPid.value, cfgCertainty.value);
+      pidModel = PidONNXModel<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp, cfgPid.value, cfgCertainty.value);
     }
 
     for (auto& track : tracks) {