diff --git a/DataFormats/Detectors/TPC/include/DataFormatsTPC/ClusterNative.h b/DataFormats/Detectors/TPC/include/DataFormatsTPC/ClusterNative.h
index fb81afdf67587..dd57fae5e0e19 100644
--- a/DataFormats/Detectors/TPC/include/DataFormatsTPC/ClusterNative.h
+++ b/DataFormats/Detectors/TPC/include/DataFormatsTPC/ClusterNative.h
@@ -291,10 +291,42 @@ struct ClusterNative {
   }
 };
 
+struct ClusterNativeNNDirection {
+  static constexpr float scalePacked = 2048.f;
+  static constexpr float maxDirection = 32767.f / scalePacked;
+  static constexpr uint16_t invalidPacked = 0;
+
+  uint16_t dydxPacked = invalidPacked;
+  uint16_t dzdxPacked = invalidPacked;
+
+  GPUd() static uint16_t pack(float v)
+  {
+    v = v < -maxDirection ? -maxDirection : (v > maxDirection ? maxDirection : v);
+    const int32_t packedSigned = static_cast<int32_t>(v * scalePacked + (v >= 0.f ? 0.5f : -0.5f));
+    return static_cast<uint16_t>(packedSigned + 32768);
+  }
+  GPUd() static float unpack(uint16_t v) { return static_cast<float>(static_cast<int32_t>(v) - 32768) * (1.f / scalePacked); }
+
+  GPUd() bool hasDirection() const { return dydxPacked != invalidPacked && dzdxPacked != invalidPacked; }
+  GPUd() float getDydx() const { return unpack(dydxPacked); }
+  GPUd() float getDzdx() const { return unpack(dzdxPacked); }
+  GPUd() void set(float dydx, float dzdx)
+  {
+    dydxPacked = pack(dydx);
+    dzdxPacked = pack(dzdx);
+  }
+  GPUd() void clear()
+  {
+    dydxPacked = invalidPacked;
+    dzdxPacked = invalidPacked;
+  }
+};
+
 // This is an index struct to access TPC clusters inside sectors and rows. It shall not own the data, but just point to
 // the data inside a buffer.
 struct ClusterNativeAccess {
   const ClusterNative* clustersLinear;
+  const ClusterNativeNNDirection* clustersLinearNNDirection = nullptr;
   const ClusterNative* clusters[constants::MAXSECTOR][constants::MAXGLOBALPADROW];
   const o2::dataformats::ConstMCTruthContainerView<o2::MCCompLabel>* clustersMCTruth;
   unsigned int nClusters[constants::MAXSECTOR][constants::MAXGLOBALPADROW];
diff --git a/GPU/GPUTracking/Base/GPUConstantMem.h b/GPU/GPUTracking/Base/GPUConstantMem.h
index 14c388e450d73..8e8c9d9be46c8 100644
--- a/GPU/GPUTracking/Base/GPUConstantMem.h
+++ b/GPU/GPUTracking/Base/GPUConstantMem.h
@@ -111,6 +111,11 @@ GPUdi() GPUconstantref() const GPUParam& GPUProcessor::Param() const
   return GetConstantMem()->param;
 }
 
+GPUdi() const o2::tpc::ClusterNativeAccess* GPUProcessor::GetClustersNative() const
+{
+  return GetConstantMem()->ioPtrs.clustersNative;
+}
+
 GPUdi() void GPUProcessor::raiseError(uint32_t code, uint32_t param1, uint32_t param2, uint32_t param3) const
 {
   GetConstantMem()->errorCodes.raiseError(code, param1, param2, param3);
diff --git a/GPU/GPUTracking/Base/GPUProcessor.h b/GPU/GPUTracking/Base/GPUProcessor.h
index 337ecfc61f79d..512d3e1d03807 100644
--- a/GPU/GPUTracking/Base/GPUProcessor.h
+++ b/GPU/GPUTracking/Base/GPUProcessor.h
@@ -23,6 +23,11 @@
 #include <algorithm>
 #endif
 
+namespace o2::tpc
+{
+struct ClusterNativeAccess;
+}
+
 namespace o2::gpu
 {
 struct GPUTrackingInOutPointers;
@@ -49,6 +54,7 @@ class GPUProcessor
 #endif
 
   GPUd() GPUconstantref() const GPUConstantMem* GetConstantMem() const; // Body in GPUConstantMem.h to avoid circular headers
+  GPUd() const o2::tpc::ClusterNativeAccess* GetClustersNative() const; // ...
   GPUd() GPUconstantref() const GPUParam& Param() const;                // ...
   GPUd() void raiseError(uint32_t code, uint32_t param1 = 0, uint32_t param2 = 0, uint32_t param3 = 0) const;
   const GPUReconstruction& GetRec() const { return *mRec; }
diff --git a/GPU/GPUTracking/Definitions/GPUSettingsList.h b/GPU/GPUTracking/Definitions/GPUSettingsList.h
index 40f7a34e699c9..dc3869279f16d 100644
--- a/GPU/GPUTracking/Definitions/GPUSettingsList.h
+++ b/GPU/GPUTracking/Definitions/GPUSettingsList.h
@@ -167,6 +167,7 @@ AddOptionRTC(dEdxClusterRejectionFlagMask, int8_t, o2::gpu::GPUTPCGMMergedTrackH
 AddOptionRTC(dEdxClusterRejectionFlagMaskAlt, int8_t, o2::gpu::GPUTPCGMMergedTrackHit::flagEdge, "", 0, "OR mask of TPC flags that will reject the cluster in alternative dEdx")
 AddOptionRTC(rejectEdgeClustersInSeeding, int8_t, 0, "", 0, "Reject edge clusters based on uncorrected track Y during seeding")
 AddOptionRTC(rejectEdgeClustersInTrackFit, int8_t, 0, "", 0, "Reject edge clusters based on uncorrected track Y during track fit")
+AddOptionRTC(useNNClusterDirection, int8_t, 0, "", 0, "Use TPC NN cluster direction estimate in track seeding")
 AddOptionRTC(tubeExtraProtectMinRow, uint8_t, 20, "", 0, "Increase Protection, decrease removal by factor 2, when below this row")
 AddOptionRTC(tubeExtraProtectEdgePads, uint8_t, 2, "", 0, "Increase Protection, decrease removal by factor 2, when on this number of pads from the edge")
 
@@ -295,7 +296,7 @@ AddOption(nnLoadFromCCDB, int, 0, "", 0, "If 1 networks are fetched from ccdb, e
 AddOption(nnCCDBDumpToFile, int, 0, "", 0, "If 1, additionally dump fetched CCDB networks to nnLocalFolder")
 AddOption(nnLocalFolder, std::string, ".", "", 0, "Local folder in which the networks will be fetched")
 AddOption(nnCCDBPath, std::string, "Users/c/csonnabe/TPC/Clusterization", "", 0, "Folder path containing the networks")
-AddOption(nnCCDBWithMomentum, std::string, "", "", 0, "Distinguishes between the network with and without momentum output for the regression")
+AddOption(nnCCDBWithMomentum, std::string, "0", "", 0, "Distinguishes between the network with and without momentum output for the regression")
 AddOption(nnCCDBClassificationLayerType, std::string, "FC", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN")
 AddOption(nnCCDBRegressionLayerType, std::string, "FC", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN")
 AddOption(nnCCDBBeamType, std::string, "pp", "", 0, "Distinguishes between networks trained for different beam types. Options: pp, pPb, PbPb")
diff --git a/GPU/GPUTracking/Global/GPUChainTracking.h b/GPU/GPUTracking/Global/GPUChainTracking.h
index 3c3531530372a..3d5666c0d53fe 100644
--- a/GPU/GPUTracking/Global/GPUChainTracking.h
+++ b/GPU/GPUTracking/Global/GPUChainTracking.h
@@ -41,6 +41,7 @@ namespace o2::tpc
 {
 struct ClusterNativeAccess;
 struct ClusterNative;
+struct ClusterNativeNNDirection;
 class CalibdEdxContainer;
 } // namespace o2::tpc
 
@@ -313,7 +314,7 @@ class GPUChainTracking : public GPUChain
   void RunTPCClusterFilter(o2::tpc::ClusterNativeAccess* clusters, std::function<o2::tpc::ClusterNative*(size_t)> allocator, bool applyClusterCuts);
   bool NeedTPCClustersOnGPU();
   void WriteReducedClusters();
-  void SortClusters(bool buildNativeGPU, bool propagateMCLabels, o2::tpc::ClusterNativeAccess* clusterAccess, o2::tpc::ClusterNative* clusters);
+  void SortClusters(bool buildNativeGPU, bool propagateMCLabels, o2::tpc::ClusterNativeAccess* clusterAccess, o2::tpc::ClusterNative* clusters, o2::tpc::ClusterNativeNNDirection* directions);
   template <int32_t I>
   int32_t RunTRDTrackingInternal();
   uint32_t StreamForSector(uint32_t sector) const;
diff --git a/GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx b/GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx
index 462b7798ce337..fc62bad7a6106 100644
--- a/GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx
+++ b/GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx
@@ -936,11 +936,16 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
   ClusterNativeAccess* tmpNativeAccess = mClusterNativeAccess.get();
   ClusterNative* tmpNativeClusters = nullptr;
   std::unique_ptr<ClusterNative[]> tmpNativeClusterBuffer;
+  ClusterNativeNNDirection* tmpNativeNNDirections = nullptr;
+  std::unique_ptr<ClusterNativeNNDirection[]> tmpNativeNNDirectionBuffer;
 
   const bool buildNativeGPU = doGPU && NeedTPCClustersOnGPU();
   const bool buildNativeHost = (mRec->GetRecoStepsOutputs() & gpudatatypes::InOutType::TPCClusters) || GetProcessingSettings().deterministicGPUReconstruction; // TODO: Should do this also when clusters are needed for later steps on the host but not requested as output
+  const bool buildNativeNNDirection = rec()->GetParam().rec.tpc.useNNClusterDirection && GetProcessingSettings().nn.applyNNclusterizer;
   const bool propagateMCLabels = buildNativeHost && GetProcessingSettings().runMC && processors()->ioPtrs.tpcPackedDigits && processors()->ioPtrs.tpcPackedDigits->tpcDigitsMC;
   const bool sortClusters = buildNativeHost && (GetProcessingSettings().deterministicGPUReconstruction || GetProcessingSettings().debugLevel >= 4);
+  const bool buildNativeNNDirectionHost = buildNativeNNDirection && (!buildNativeGPU || sortClusters);
+  bool buildNativeNNDirectionGPU = buildNativeNNDirection;
 
   if (GetProcessingSettings().runMC && (!processors()->ioPtrs.tpcPackedDigits || !processors()->ioPtrs.tpcPackedDigits->tpcDigitsMC)) {
     GPUWarning("Requested to process MC labels, but no labels present");
@@ -951,6 +956,9 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
   mInputsHost->mNClusterNative = mInputsShadow->mNClusterNative = mRec->MemoryScalers()->nTPCHits * tpcHitLowOccupancyScalingFactor;
   if (buildNativeGPU) {
     AllocateRegisteredMemory(mInputsHost->mResourceClusterNativeBuffer);
+    if (buildNativeNNDirection) {
+      AllocateRegisteredMemory(mInputsHost->mResourceClusterNativeNNDirectionBuffer);
+    }
   }
   if (mWaitForFinalInputs && GetProcessingSettings().nTPCClustererLanes > 6) {
     GPUFatal("ERROR, mWaitForFinalInputs cannot be called with nTPCClustererLanes > 6");
@@ -962,9 +970,17 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
     if (!GetProcessingSettings().tpcApplyClusterFilterOnCPU) {
       AllocateRegisteredMemory(mInputsHost->mResourceClusterNativeOutput, GetProcessingSettings().tpcWriteClustersAfterRejection ? nullptr : mSubOutputControls[GPUTrackingOutputs::getIndex(&GPUTrackingOutputs::clustersNative)]);
       tmpNativeClusters = mInputsHost->mPclusterNativeOutput;
+      if (buildNativeNNDirectionHost) {
+        AllocateRegisteredMemory(mInputsHost->mResourceClusterNativeNNDirectionOutput);
+        tmpNativeNNDirections = mInputsHost->mPclusterNativeNNDirectionOutput;
+      }
     } else {
       tmpNativeClusterBuffer = std::make_unique<ClusterNative[]>(mInputsHost->mNClusterNative);
       tmpNativeClusters = tmpNativeClusterBuffer.get();
+      if (buildNativeNNDirectionHost) {
+        tmpNativeNNDirectionBuffer = std::make_unique<ClusterNativeNNDirection[]>(mInputsHost->mNClusterNative);
+        tmpNativeNNDirections = tmpNativeNNDirectionBuffer.get();
+      }
     }
   }
 
@@ -1216,6 +1232,9 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
             transferRunning[lane] = 2;
           }
           runKernel<GPUMemClean16>({GetGridAutoStep(lane, RecoStep::TPCClusterFinding), krnlRunRangeNone, {nullptr, waitEvent}}, clustererShadow.mPclusterInRow, GPUTPCGeometry::NROWS * sizeof(*clustererShadow.mPclusterInRow));
+          if (clustererShadow.mPclusterNNDirectionByRow != nullptr) {
+            runKernel<GPUMemClean16>({GetGridAutoStep(lane, RecoStep::TPCClusterFinding), krnlRunRangeNone}, clustererShadow.mPclusterNNDirectionByRow, GPUTPCGeometry::NROWS * clusterer.mNMaxClusterPerRow * sizeof(*clustererShadow.mPclusterNNDirectionByRow));
+          }
         }
 
         const auto nRegularClusters = clusterer.mPmemory->counters.nClusters;
@@ -1438,8 +1457,14 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
             if (!GetProcessingSettings().tpccfGatherKernel) {
               GPUMemCpyAlways(RecoStep::TPCClusterFinding, (void*)&mInputsShadow->mPclusterNativeBuffer[nClsTotal], (const void*)&clustererShadow.mPclusterByRow[j * clusterer.mNMaxClusterPerRow], sizeof(mIOPtrs.clustersNative->clustersLinear[0]) * clusterer.mPclusterInRow[j], mRec->NStreams() - 1, -2);
             }
+            if (buildNativeNNDirection && clustererShadow.mPclusterNNDirectionByRow != nullptr) {
+              GPUMemCpyAlways(RecoStep::TPCClusterFinding, (void*)&mInputsShadow->mPclusterNativeNNDirectionBuffer[nClsTotal], (const void*)&clustererShadow.mPclusterNNDirectionByRow[j * clusterer.mNMaxClusterPerRow], sizeof(mInputsShadow->mPclusterNativeNNDirectionBuffer[0]) * clusterer.mPclusterInRow[j], mRec->NStreams() - 1, -2);
+            }
           } else if (buildNativeHost) {
             GPUMemCpyAlways(RecoStep::TPCClusterFinding, (void*)&tmpNativeClusters[nClsTotal], (const void*)&clustererShadow.mPclusterByRow[j * clusterer.mNMaxClusterPerRow], sizeof(mIOPtrs.clustersNative->clustersLinear[0]) * clusterer.mPclusterInRow[j], mRec->NStreams() - 1, false);
+            if (buildNativeNNDirectionHost && clustererShadow.mPclusterNNDirectionByRow != nullptr) {
+              GPUMemCpyAlways(RecoStep::TPCClusterFinding, (void*)&tmpNativeNNDirections[nClsTotal], (const void*)&clustererShadow.mPclusterNNDirectionByRow[j * clusterer.mNMaxClusterPerRow], sizeof(tmpNativeNNDirections[0]) * clusterer.mPclusterInRow[j], mRec->NStreams() - 1, false);
+            }
           }
           tmpNativeAccess->nClusters[iSector][j] += clusterer.mPclusterInRow[j];
           nClsTotal += clusterer.mPclusterInRow[j];
@@ -1468,6 +1493,9 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
         mOutputQueue.emplace_back(outputQueueEntry{(void*)((char*)&tmpNativeClusters[nClsFirst] - (char*)&tmpNativeClusters[0]), &mInputsShadow->mPclusterNativeBuffer[nClsFirst], (nClsTotal - nClsFirst) * sizeof(tmpNativeClusters[0]), RecoStep::TPCClusterFinding});
       } else {
         GPUMemCpy(RecoStep::TPCClusterFinding, (void*)&tmpNativeClusters[nClsFirst], (const void*)&mInputsShadow->mPclusterNativeBuffer[nClsFirst], (nClsTotal - nClsFirst) * sizeof(tmpNativeClusters[0]), mRec->NStreams() - 1, false);
+        if (buildNativeNNDirectionHost) {
+          GPUMemCpy(RecoStep::TPCClusterFinding, (void*)&tmpNativeNNDirections[nClsFirst], (const void*)&mInputsShadow->mPclusterNativeNNDirectionBuffer[nClsFirst], (nClsTotal - nClsFirst) * sizeof(tmpNativeNNDirections[0]), mRec->NStreams() - 1, false);
+        }
       }
     }
 
@@ -1541,6 +1569,11 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
     mInputsHost->mNClusterNative = mInputsShadow->mNClusterNative = nClsTotal;
     AllocateRegisteredMemory(mInputsHost->mResourceClusterNativeOutput, GetProcessingSettings().tpcWriteClustersAfterRejection ? nullptr : mSubOutputControls[GPUTrackingOutputs::getIndex(&GPUTrackingOutputs::clustersNative)]);
     tmpNativeClusters = mInputsHost->mPclusterNativeOutput;
+    if (buildNativeNNDirectionHost) {
+      AllocateRegisteredMemory(mInputsHost->mResourceClusterNativeNNDirectionOutput);
+      tmpNativeNNDirections = mInputsHost->mPclusterNativeNNDirectionOutput;
+      GPUMemCpy(RecoStep::TPCClusterFinding, (void*)tmpNativeNNDirections, (const void*)mInputsShadow->mPclusterNativeNNDirectionBuffer, nClsTotal * sizeof(tmpNativeNNDirections[0]), mRec->NStreams() - 1, false);
+    }
     for (uint32_t i = outputQueueStart; i < mOutputQueue.size(); i++) {
       mOutputQueue[i].dst = (char*)tmpNativeClusters + (size_t)mOutputQueue[i].dst;
     }
@@ -1548,6 +1581,7 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
 
   if (buildNativeHost) {
     tmpNativeAccess->clustersLinear = tmpNativeClusters;
+    tmpNativeAccess->clustersLinearNNDirection = tmpNativeNNDirections;
     tmpNativeAccess->clustersMCTruth = mcLabelsConstView;
     tmpNativeAccess->setOffsetPtrs();
     mIOPtrs.clustersNative = tmpNativeAccess;
@@ -1559,6 +1593,9 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
       };
       RunTPCClusterFilter(tmpNativeAccess, allocator, false);
       nClsTotal = tmpNativeAccess->nClustersTotal;
+      tmpNativeAccess->clustersLinearNNDirection = nullptr;
+      tmpNativeNNDirections = nullptr;
+      buildNativeNNDirectionGPU = false;
     }
   }
 
@@ -1571,6 +1608,7 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
     WriteToConstantMemory(RecoStep::TPCClusterFinding, (char*)&processors()->ioPtrs - (char*)processors(), &processorsShadow()->ioPtrs, sizeof(processorsShadow()->ioPtrs), 0);
     *mInputsHost->mPclusterNativeAccess = *mIOPtrs.clustersNative;
     mInputsHost->mPclusterNativeAccess->clustersLinear = mInputsShadow->mPclusterNativeBuffer;
+    mInputsHost->mPclusterNativeAccess->clustersLinearNNDirection = buildNativeNNDirectionGPU ? mInputsShadow->mPclusterNativeNNDirectionBuffer : nullptr;
     mInputsHost->mPclusterNativeAccess->setOffsetPtrs();
     TransferMemoryResourceLinkToGPU(RecoStep::TPCClusterFinding, mInputsHost->mResourceClusterNativeAccess, 0);
   }
@@ -1581,7 +1619,7 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
     SynchronizeStream(0);
   }
   if (sortClusters) {
-    SortClusters(buildNativeGPU, propagateMCLabels, tmpNativeAccess, tmpNativeClusters);
+    SortClusters(buildNativeGPU, propagateMCLabels, tmpNativeAccess, tmpNativeClusters, tmpNativeNNDirections);
   }
   mRec->MemoryScalers()->nTPCHits = nClsTotal;
   mRec->PopNonPersistentMemory(RecoStep::TPCClusterFinding, qStr2Tag("TPCCLUST"));
@@ -1598,12 +1636,13 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
   return 0;
 }
 
-void GPUChainTracking::SortClusters(bool buildNativeGPU, bool propagateMCLabels, ClusterNativeAccess* clusterAccess, ClusterNative* clusters)
+void GPUChainTracking::SortClusters(bool buildNativeGPU, bool propagateMCLabels, ClusterNativeAccess* clusterAccess, ClusterNative* clusters, ClusterNativeNNDirection* directions)
 {
-  if (propagateMCLabels) {
+  if (propagateMCLabels || directions) {
     std::vector<uint32_t> clsOrder(clusterAccess->nClustersTotal);
     std::iota(clsOrder.begin(), clsOrder.end(), 0);
     std::vector<ClusterNative> tmpClusters;
+    std::vector<ClusterNativeNNDirection> tmpDirections;
     for (uint32_t i = 0; i < NSECTORS; i++) {
       for (uint32_t j = 0; j < GPUTPCGeometry::NROWS; j++) {
         const uint32_t offset = clusterAccess->clusterOffset[i][j];
@@ -1615,34 +1654,44 @@ void GPUChainTracking::SortClusters(bool buildNativeGPU, bool propagateMCLabels,
         for (uint32_t k = 0; k < tmpClusters.size(); k++) {
           clusters[offset + k] = tmpClusters[clsOrder[offset + k] - offset];
         }
+        if (directions) {
+          tmpDirections.resize(clusterAccess->nClusters[i][j]);
+          memcpy(tmpDirections.data(), &directions[offset], clusterAccess->nClusters[i][j] * sizeof(tmpDirections[0]));
+          for (uint32_t k = 0; k < tmpDirections.size(); k++) {
+            directions[offset + k] = tmpDirections[clsOrder[offset + k] - offset];
+          }
+        }
       }
     }
     tmpClusters.clear();
+    tmpDirections.clear();
+
+    if (propagateMCLabels) {
+      std::pair<o2::dataformats::ConstMCLabelContainer*, o2::dataformats::ConstMCLabelContainerView*> labelBuffer;
+      GPUOutputControl* labelOutput = mSubOutputControls[GPUTrackingOutputs::getIndex(&GPUTrackingOutputs::clusterLabels)];
+      std::unique_ptr<ConstMCLabelContainerView> tmpUniqueContainerView;
+      std::unique_ptr<ConstMCLabelContainer> tmpUniqueContainerBuffer;
+      if (labelOutput && labelOutput->allocator) {
+        ClusterNativeAccess::ConstMCLabelContainerViewWithBuffer* labelContainer = reinterpret_cast<ClusterNativeAccess::ConstMCLabelContainerViewWithBuffer*>(labelOutput->allocator(0));
+        labelBuffer = {&labelContainer->first, &labelContainer->second};
+      } else {
+        tmpUniqueContainerView = std::move(mIOMem.clusterNativeMCView);
+        tmpUniqueContainerBuffer = std::move(mIOMem.clusterNativeMCBuffer);
+        mIOMem.clusterNativeMCView = std::make_unique<ConstMCLabelContainerView>();
+        mIOMem.clusterNativeMCBuffer = std::make_unique<ConstMCLabelContainer>();
+        labelBuffer = {mIOMem.clusterNativeMCBuffer.get(), mIOMem.clusterNativeMCView.get()};
+      }
 
-    std::pair<o2::dataformats::ConstMCLabelContainer*, o2::dataformats::ConstMCLabelContainerView*> labelBuffer;
-    GPUOutputControl* labelOutput = mSubOutputControls[GPUTrackingOutputs::getIndex(&GPUTrackingOutputs::clusterLabels)];
-    std::unique_ptr<ConstMCLabelContainerView> tmpUniqueContainerView;
-    std::unique_ptr<ConstMCLabelContainer> tmpUniqueContainerBuffer;
-    if (labelOutput && labelOutput->allocator) {
-      ClusterNativeAccess::ConstMCLabelContainerViewWithBuffer* labelContainer = reinterpret_cast<ClusterNativeAccess::ConstMCLabelContainerViewWithBuffer*>(labelOutput->allocator(0));
-      labelBuffer = {&labelContainer->first, &labelContainer->second};
-    } else {
-      tmpUniqueContainerView = std::move(mIOMem.clusterNativeMCView);
-      tmpUniqueContainerBuffer = std::move(mIOMem.clusterNativeMCBuffer);
-      mIOMem.clusterNativeMCView = std::make_unique<ConstMCLabelContainerView>();
-      mIOMem.clusterNativeMCBuffer = std::make_unique<ConstMCLabelContainer>();
-      labelBuffer = {mIOMem.clusterNativeMCBuffer.get(), mIOMem.clusterNativeMCView.get()};
-    }
-
-    o2::dataformats::MCLabelContainer tmpContainer;
-    for (uint32_t i = 0; i < clusterAccess->nClustersTotal; i++) {
-      for (const auto& element : clusterAccess->clustersMCTruth->getLabels(clsOrder[i])) {
-        tmpContainer.addElement(i, element);
+      o2::dataformats::MCLabelContainer tmpContainer;
+      for (uint32_t i = 0; i < clusterAccess->nClustersTotal; i++) {
+        for (const auto& element : clusterAccess->clustersMCTruth->getLabels(clsOrder[i])) {
+          tmpContainer.addElement(i, element);
+        }
       }
+      tmpContainer.flatten_to(*labelBuffer.first);
+      *labelBuffer.second = *labelBuffer.first;
+      clusterAccess->clustersMCTruth = labelBuffer.second;
     }
-    tmpContainer.flatten_to(*labelBuffer.first);
-    *labelBuffer.second = *labelBuffer.first;
-    clusterAccess->clustersMCTruth = labelBuffer.second;
   } else {
     for (uint32_t i = 0; i < NSECTORS; i++) {
       for (uint32_t j = 0; j < GPUTPCGeometry::NROWS; j++) {
@@ -1652,5 +1701,8 @@ void GPUChainTracking::SortClusters(bool buildNativeGPU, bool propagateMCLabels,
   }
   if (buildNativeGPU) {
     GPUMemCpy(RecoStep::TPCClusterFinding, (void*)mInputsShadow->mPclusterNativeBuffer, (const void*)clusters, clusterAccess->nClustersTotal * sizeof(clusters[0]), -1, true);
+    if (directions) {
+      GPUMemCpy(RecoStep::TPCClusterFinding, (void*)mInputsShadow->mPclusterNativeNNDirectionBuffer, (const void*)directions, clusterAccess->nClustersTotal * sizeof(directions[0]), -1, true);
+    }
   }
 }
diff --git a/GPU/GPUTracking/Global/GPUChainTrackingTransformation.cxx b/GPU/GPUTracking/Global/GPUChainTrackingTransformation.cxx
index 46bcd6931c302..4c9596476fc03 100644
--- a/GPU/GPUTracking/Global/GPUChainTrackingTransformation.cxx
+++ b/GPU/GPUTracking/Global/GPUChainTrackingTransformation.cxx
@@ -44,12 +44,19 @@ int32_t GPUChainTracking::ConvertNativeToClusterData()
   if (mRec->IsGPU() && !(mRec->GetRecoStepsGPU() & gpudatatypes::RecoStep::TPCClusterFinding) && NeedTPCClustersOnGPU()) {
     mInputsHost->mNClusterNative = mInputsShadow->mNClusterNative = mIOPtrs.clustersNative->nClustersTotal;
     AllocateRegisteredMemory(mInputsHost->mResourceClusterNativeBuffer);
+    if (mIOPtrs.clustersNative->clustersLinearNNDirection && param().rec.tpc.useNNClusterDirection) {
+      AllocateRegisteredMemory(mInputsHost->mResourceClusterNativeNNDirectionBuffer);
+    }
     processorsShadow()->ioPtrs.clustersNative = mInputsShadow->mPclusterNativeAccess;
     WriteToConstantMemory(RecoStep::TPCConversion, (char*)&processors()->ioPtrs - (char*)processors(), &processorsShadow()->ioPtrs, sizeof(processorsShadow()->ioPtrs), 0);
     *mInputsHost->mPclusterNativeAccess = *mIOPtrs.clustersNative;
     mInputsHost->mPclusterNativeAccess->clustersLinear = mInputsShadow->mPclusterNativeBuffer;
+    mInputsHost->mPclusterNativeAccess->clustersLinearNNDirection = (mIOPtrs.clustersNative->clustersLinearNNDirection && param().rec.tpc.useNNClusterDirection) ? mInputsShadow->mPclusterNativeNNDirectionBuffer : nullptr;
     mInputsHost->mPclusterNativeAccess->setOffsetPtrs();
     GPUMemCpy(RecoStep::TPCConversion, mInputsShadow->mPclusterNativeBuffer, mIOPtrs.clustersNative->clustersLinear, sizeof(mIOPtrs.clustersNative->clustersLinear[0]) * mIOPtrs.clustersNative->nClustersTotal, 0, true);
+    if (mInputsHost->mPclusterNativeAccess->clustersLinearNNDirection) {
+      GPUMemCpy(RecoStep::TPCConversion, mInputsShadow->mPclusterNativeNNDirectionBuffer, mIOPtrs.clustersNative->clustersLinearNNDirection, sizeof(mIOPtrs.clustersNative->clustersLinearNNDirection[0]) * mIOPtrs.clustersNative->nClustersTotal, 0, true);
+    }
     TransferMemoryResourceLinkToGPU(RecoStep::TPCConversion, mInputsHost->mResourceClusterNativeAccess, 0);
     transferClusters = true;
   }
diff --git a/GPU/GPUTracking/Global/GPUTrackingInputProvider.cxx b/GPU/GPUTracking/Global/GPUTrackingInputProvider.cxx
index dc47b6c0a6663..236c19d7c5ebf 100644
--- a/GPU/GPUTracking/Global/GPUTrackingInputProvider.cxx
+++ b/GPU/GPUTracking/Global/GPUTrackingInputProvider.cxx
@@ -52,6 +52,14 @@ void* GPUTrackingInputProvider::SetPointersInputClusterNativeBuffer(void* mem)
   return mem;
 }
 
+void* GPUTrackingInputProvider::SetPointersInputClusterNativeNNDirectionBuffer(void* mem)
+{
+  if (mHoldTPCClusterNative && mRec->GetParam().rec.tpc.useNNClusterDirection) {
+    computePointerWithAlignment(mem, mPclusterNativeNNDirectionBuffer, mNClusterNative);
+  }
+  return mem;
+}
+
 void* GPUTrackingInputProvider::SetPointersInputClusterNativeOutput(void* mem)
 {
   if (mHoldTPCClusterNativeOutput) {
@@ -60,6 +68,14 @@ void* GPUTrackingInputProvider::SetPointersInputClusterNativeOutput(void* mem)
   return mem;
 }
 
+void* GPUTrackingInputProvider::SetPointersInputClusterNativeNNDirectionOutput(void* mem)
+{
+  if (mHoldTPCClusterNativeOutput && mRec->GetParam().rec.tpc.useNNClusterDirection) {
+    computePointerWithAlignment(mem, mPclusterNativeNNDirectionOutput, mNClusterNative);
+  }
+  return mem;
+}
+
 void* GPUTrackingInputProvider::SetPointersErrorCodes(void* mem)
 {
   computePointerWithAlignment(mem, mErrorCodes, 4 * GPUErrors::getMaxErrors() + 1);
@@ -95,7 +111,9 @@ void GPUTrackingInputProvider::RegisterMemoryAllocation()
   mResourceOccupancyMap = mRec->RegisterMemoryAllocation(this, &GPUTrackingInputProvider::SetPointersTPCOccupancyMap, GPUMemoryResource::MEMORY_INOUT | GPUMemoryResource::MEMORY_CUSTOM, "OccupancyMap");
   mResourceClusterNativeAccess = mRec->RegisterMemoryAllocation(this, &GPUTrackingInputProvider::SetPointersInputClusterNativeAccess, GPUMemoryResource::MEMORY_INPUT, "ClusterNativeAccess");
   mResourceClusterNativeBuffer = mRec->RegisterMemoryAllocation(this, &GPUTrackingInputProvider::SetPointersInputClusterNativeBuffer, GPUMemoryResource::MEMORY_INPUT_FLAG | GPUMemoryResource::MEMORY_GPU | GPUMemoryResource::MEMORY_EXTERNAL | GPUMemoryResource::MEMORY_CUSTOM, "ClusterNativeBuffer");
+  mResourceClusterNativeNNDirectionBuffer = mRec->RegisterMemoryAllocation(this, &GPUTrackingInputProvider::SetPointersInputClusterNativeNNDirectionBuffer, GPUMemoryResource::MEMORY_INPUT_FLAG | GPUMemoryResource::MEMORY_GPU | GPUMemoryResource::MEMORY_EXTERNAL | GPUMemoryResource::MEMORY_CUSTOM, "ClusterNativeNNDirectionBuffer");
   mResourceClusterNativeOutput = mRec->RegisterMemoryAllocation(this, &GPUTrackingInputProvider::SetPointersInputClusterNativeOutput, GPUMemoryResource::MEMORY_OUTPUT_FLAG | GPUMemoryResource::MEMORY_HOST | GPUMemoryResource::MEMORY_CUSTOM, "ClusterNativeOutput");
+  mResourceClusterNativeNNDirectionOutput = mRec->RegisterMemoryAllocation(this, &GPUTrackingInputProvider::SetPointersInputClusterNativeNNDirectionOutput, GPUMemoryResource::MEMORY_OUTPUT_FLAG | GPUMemoryResource::MEMORY_HOST | GPUMemoryResource::MEMORY_CUSTOM, "ClusterNativeNNDirectionOutput");
   mResourceTRD = mRec->RegisterMemoryAllocation(this, &GPUTrackingInputProvider::SetPointersInputTRD, GPUMemoryResource::MEMORY_INPUT_FLAG | GPUMemoryResource::MEMORY_GPU | GPUMemoryResource::MEMORY_EXTERNAL | GPUMemoryResource::MEMORY_CUSTOM, "TRDInputBuffer");
 }
 
diff --git a/GPU/GPUTracking/Global/GPUTrackingInputProvider.h b/GPU/GPUTracking/Global/GPUTrackingInputProvider.h
index 7aee803a03ace..017934a8a676b 100644
--- a/GPU/GPUTracking/Global/GPUTrackingInputProvider.h
+++ b/GPU/GPUTracking/Global/GPUTrackingInputProvider.h
@@ -21,6 +21,7 @@
 namespace o2::tpc
 {
 struct ClusterNative;
+struct ClusterNativeNNDirection;
 struct ClusterNativeAccess;
 } // namespace o2::tpc
 
@@ -44,7 +45,9 @@ class GPUTrackingInputProvider : public GPUProcessor
   void* SetPointersInputZS(void* mem);
   void* SetPointersInputClusterNativeAccess(void* mem);
   void* SetPointersInputClusterNativeBuffer(void* mem);
+  void* SetPointersInputClusterNativeNNDirectionBuffer(void* mem);
   void* SetPointersInputClusterNativeOutput(void* mem);
+  void* SetPointersInputClusterNativeNNDirectionOutput(void* mem);
   void* SetPointersInputTRD(void* mem);
   void* SetPointersErrorCodes(void* mem);
 #endif
@@ -52,7 +55,9 @@ class GPUTrackingInputProvider : public GPUProcessor
   uint16_t mResourceZS = -1;
   uint16_t mResourceClusterNativeAccess = -1;
   uint16_t mResourceClusterNativeBuffer = -1;
+  uint16_t mResourceClusterNativeNNDirectionBuffer = -1;
   uint16_t mResourceClusterNativeOutput = -1;
+  uint16_t mResourceClusterNativeNNDirectionOutput = -1;
   uint16_t mResourceErrorCodes = -1;
   uint16_t mResourceTRD = -1;
   uint16_t mResourceOccupancyMap = -1;
@@ -78,7 +83,9 @@ class GPUTrackingInputProvider : public GPUProcessor
 
   o2::tpc::ClusterNativeAccess* mPclusterNativeAccess = nullptr;
   o2::tpc::ClusterNative* mPclusterNativeBuffer = nullptr;
+  o2::tpc::ClusterNativeNNDirection* mPclusterNativeNNDirectionBuffer = nullptr;
   o2::tpc::ClusterNative* mPclusterNativeOutput = nullptr;
+  o2::tpc::ClusterNativeNNDirection* mPclusterNativeNNDirectionOutput = nullptr;
 
   uint32_t* mTPCClusterOccupancyMap = nullptr;
   uint32_t mTPCClusterOccupancyMapSize = 0;
diff --git a/GPU/GPUTracking/SectorTracker/GPUTPCNeighboursFinder.cxx b/GPU/GPUTracking/SectorTracker/GPUTPCNeighboursFinder.cxx
index 5151427377b05..1ba0502c0839f 100644
--- a/GPU/GPUTracking/SectorTracker/GPUTPCNeighboursFinder.cxx
+++ b/GPU/GPUTracking/SectorTracker/GPUTPCNeighboursFinder.cxx
@@ -120,6 +120,8 @@ GPUdii() void GPUTPCNeighboursFinder::Thread<0>(int32_t /*nBlocks*/, int32_t nTh
     const GPUglobalref() cahit2& hitData = pHitData[lHitNumberOffset + ih];
     const float y = y0 + hitData.x * stepY;
     const float z = z0 + hitData.y * stepZ;
+    float nnDydx = 0.f, nnDzdx = 0.f;
+    const bool useNNDir = tracker.Param().rec.tpc.useNNClusterDirection && tracker.HitNNDirection(row, ih, nnDydx, nnDzdx) && CAMath::Abs(nnDydx) < 10.f && CAMath::Abs(nnDzdx) < 10.f;
 
     uint32_t nNeighUp = 0;
     float minZ, maxZ, minY, maxY;
@@ -127,8 +129,8 @@ GPUdii() void GPUTPCNeighboursFinder::Thread<0>(int32_t /*nBlocks*/, int32_t nTh
     int32_t nY;
 
     { // area in the upper row
-      const float yy = y * s.mUpTx;
-      const float zz = z * kAreaSlopeZUp;
+      const float yy = useNNDir ? y + nnDydx * s.mUpDx : y * s.mUpTx;
+      const float zz = useNNDir ? z + nnDzdx * s.mUpDx : z * kAreaSlopeZUp;
       minZ = zz - kAreaSizeZUp;
       maxZ = zz + kAreaSizeZUp;
       minY = yy - kAreaSizeY;
@@ -196,8 +198,8 @@ GPUdii() void GPUTPCNeighboursFinder::Thread<0>(int32_t /*nBlocks*/, int32_t nTh
     }
 
     { // area in the lower row
-      const float yy = y * s.mDnTx;
-      const float zz = z * kAreaSlopeZDn;
+      const float yy = useNNDir ? y + nnDydx * s.mDnDx : y * s.mDnTx;
+      const float zz = useNNDir ? z + nnDzdx * s.mDnDx : z * kAreaSlopeZDn;
       minZ = zz - kAreaSizeZDn;
       maxZ = zz + kAreaSizeZDn;
       minY = yy - kAreaSizeY;
diff --git a/GPU/GPUTracking/SectorTracker/GPUTPCTracker.h b/GPU/GPUTracking/SectorTracker/GPUTPCTracker.h
index cb0c8d4a76fdb..12e3bddfb796e 100644
--- a/GPU/GPUTracking/SectorTracker/GPUTPCTracker.h
+++ b/GPU/GPUTracking/SectorTracker/GPUTPCTracker.h
@@ -26,6 +26,7 @@
 #include "GPUTPCTrackParam.h"
 #include "GPUTPCTracklet.h"
 #include "GPUProcessor.h"
+#include "DataFormatsTPC/ClusterNative.h"
 
 namespace o2::gpu
 {
@@ -151,6 +152,22 @@ class GPUTPCTracker : public GPUProcessor
 
   GPUhd() int32_t HitInputID(const GPUTPCRow& row, int32_t hitIndex) const { return mData.ClusterDataIndex(row, hitIndex); }
 
+  GPUd() bool HitNNDirection(const GPUTPCRow& row, int32_t hitIndex, float& dydx, float& dzdx) const
+  {
+    const o2::tpc::ClusterNativeAccess* native = GetClustersNative();
+    if (native == nullptr || native->clustersLinearNNDirection == nullptr) {
+      return false;
+    }
+    const uint32_t sectorOffset = native->clusterOffset[mISector][0];
+    const o2::tpc::ClusterNativeNNDirection& direction = native->clustersLinearNNDirection[sectorOffset + mData.ClusterDataIndex(row, hitIndex)];
+    if (!direction.hasDirection()) {
+      return false;
+    }
+    dydx = direction.getDydx();
+    dzdx = direction.getDzdx();
+    return true;
+  }
+
   /**
    * The hit weight is used to determine whether a hit belongs to a certain tracklet or another one
    * competing for the same hit. The tracklet that has a higher weight wins. Comparison is done
diff --git a/GPU/GPUTracking/SectorTracker/GPUTPCTrackletConstructor.cxx b/GPU/GPUTracking/SectorTracker/GPUTPCTrackletConstructor.cxx
index 33a3264a87ab3..00f7c4f0dc449 100644
--- a/GPU/GPUTracking/SectorTracker/GPUTPCTrackletConstructor.cxx
+++ b/GPU/GPUTracking/SectorTracker/GPUTPCTrackletConstructor.cxx
@@ -170,9 +170,16 @@ GPUdic(2, 1) void GPUTPCTrackletConstructor::UpdateTracklet(int32_t /*nBlocks*/,
 
         float ri = 1.f / CAMath::Sqrt(dx * dx + dy * dy);
         if (iRow == (int32_t)r.mStartRow + 2) {
-          tParam.SetSinPhi(dy * ri);
+          float nnDydx = 0.f, nnDzdx = 0.f;
+          if (tracker.Param().rec.tpc.useNNClusterDirection && tracker.HitNNDirection(row, seedIH, nnDydx, nnDzdx) && CAMath::Abs(nnDydx) < 10.f && CAMath::Abs(nnDzdx) < 10.f) {
+            const float nnNormI = CAMath::InvSqrt(1.f + nnDydx * nnDydx);
+            tParam.SetSinPhi(0.5f * (dy * ri + nnDydx * nnNormI));
+            tParam.SetDzDs(0.5f * (dz * ri + nnDzdx * nnNormI));
+          } else {
+            tParam.SetSinPhi(dy * ri);
+            tParam.SetDzDs(dz * ri);
+          }
           tParam.SetSignCosPhi(dx);
-          tParam.SetDzDs(dz * ri);
           float err2Y, err2Z;
           tracker.GetErrors2Seeding(iRow, tParam, -1.f, err2Y, err2Z); // Use correct time
           tParam.SetCov(0, err2Y);
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCClusterFinder.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCClusterFinder.cxx
index 06709ef2d4a7e..269da7488032e 100644
--- a/GPU/GPUTracking/TPCClusterFinder/GPUTPCClusterFinder.cxx
+++ b/GPU/GPUTracking/TPCClusterFinder/GPUTPCClusterFinder.cxx
@@ -97,6 +97,11 @@ void* GPUTPCClusterFinder::SetPointersScratch(void* mem)
   computePointerWithAlignment(mem, mPchargeMap, TPCMapMemoryLayout<decltype(*mPchargeMap)>::items(mRec->GetProcessingSettings().overrideClusterizerFragmentLen));
   computePointerWithAlignment(mem, mPpeakMap, TPCMapMemoryLayout<decltype(*mPpeakMap)>::items(mRec->GetProcessingSettings().overrideClusterizerFragmentLen));
   computePointerWithAlignment(mem, mPclusterByRow, GPUTPCGeometry::NROWS * mNMaxClusterPerRow);
+  if (mRec->GetProcessingSettings().nn.applyNNclusterizer && mRec->GetParam().rec.tpc.useNNClusterDirection) {
+    computePointerWithAlignment(mem, mPclusterNNDirectionByRow, GPUTPCGeometry::NROWS * mNMaxClusterPerRow);
+  } else {
+    mPclusterNNDirectionByRow = nullptr;
+  }
   if ((mRec->GetRecoStepsGPU() & gpudatatypes::RecoStep::TPCClusterFinding)) {
     computePointerWithAlignment(mem, mPscanBuf, mBufSize * mNBufs);
   }
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCClusterFinder.h b/GPU/GPUTracking/TPCClusterFinder/GPUTPCClusterFinder.h
index d169440a8d972..c62c1f4d3ae18 100644
--- a/GPU/GPUTracking/TPCClusterFinder/GPUTPCClusterFinder.h
+++ b/GPU/GPUTracking/TPCClusterFinder/GPUTPCClusterFinder.h
@@ -36,6 +36,7 @@ class ConstMCTruthContainerView;
 namespace tpc
 {
 struct ClusterNative;
+struct ClusterNativeNNDirection;
 struct ClusterNativeAccess;
 class Digit;
 } // namespace tpc
@@ -114,6 +115,7 @@ class GPUTPCClusterFinder : public GPUProcessor
   uint32_t* mPindexMap = nullptr;
   uint32_t* mPclusterInRow = nullptr;
   tpc::ClusterNative* mPclusterByRow = nullptr;
+  tpc::ClusterNativeNNDirection* mPclusterNNDirectionByRow = nullptr;
   GPUTPCClusterMCInterimArray* mPlabelsByRow = nullptr;
   int32_t* mPscanBuf = nullptr;
   HIPTailDescriptor* mPhipTailsByRow = nullptr;
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h
index 7aa23489eb2f4..92c10e06600dc 100644
--- a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h
+++ b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h
@@ -54,6 +54,7 @@ class GPUTPCNNClusterizer : public GPUProcessor
   int32_t mNnClusterizerBoundaryFillValue = -1;
   int32_t mNnClusterizerModelClassNumOutputNodes = -1;
   int32_t mNnClusterizerModelReg1NumOutputNodes = -1;
+  int32_t mNnClusterizerUseMomentumVector = 0;
   int32_t mNnClusterizerModelReg2NumOutputNodes = -1;
   int32_t mNnInferenceInputDType = 0;  // 0: float16, 1: float32
   int32_t mNnInferenceOutputDType = 0; // 0: float16, 1: float32
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx
index 96b8a1d7ed2fd..6552ede836ec4 100644
--- a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx
+++ b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx
@@ -138,6 +138,9 @@ void GPUTPCNNClusterizerHost::initClusterizer(const GPUSettingsProcessingNNclust
   if (!settings.nnClusterizerUseCfRegression) {
     if (mModelClass.getNumOutputNodes()[0][1] == 1 || !mModelReg2.isInitialized()) {
       clustererNN.mNnClusterizerModelReg1NumOutputNodes = mModelReg1.getNumOutputNodes()[0][1];
+      if (clustererNN.mNnClusterizerModelReg1NumOutputNodes > 5) {
+        clustererNN.mNnClusterizerUseMomentumVector = true;
+      }
     } else {
       clustererNN.mNnClusterizerModelReg1NumOutputNodes = mModelReg1.getNumOutputNodes()[0][1];
       clustererNN.mNnClusterizerModelReg2NumOutputNodes = mModelReg2.getNumOutputNodes()[0][1];
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.cxx
index 693ee4dd78e8d..9b5ddde43611c 100644
--- a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.cxx
+++ b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.cxx
@@ -18,6 +18,7 @@
 #include "GPUConstantMem.h"
 #include "GPUTPCClusterFinder.h"
 #include "GPUTPCGeometry.h"
+#include "DataFormatsTPC/ClusterNative.h"
 
 using namespace o2::gpu;
 using namespace o2::gpu::tpccf;
@@ -38,6 +39,17 @@ using namespace o2::gpu::tpccf;
 
 static_assert(GPUTPCNNClusterizerKernels::SCRATCH_PAD_WORK_GROUP_SIZE == GPUTPCCFClusterizer::SCRATCH_PAD_WORK_GROUP_SIZE, "Work group sizes do not match");
 
+GPUd() static o2::tpc::ClusterNativeNNDirection getClass1NNDirection(const GPUTPCNNClusterizer& nn, int8_t dtype, int32_t base)
+{
+  o2::tpc::ClusterNativeNNDirection direction;
+  if (dtype == 0) {
+    direction.set(nn.mOutputDataReg1_32[base + 5], nn.mOutputDataReg1_32[base + 6]);
+  } else {
+    direction.set(nn.mOutputDataReg1_16[base + 5].ToFloat(), nn.mOutputDataReg1_16[base + 6].ToFloat());
+  }
+  return direction;
+}
+
 // Defining individual thread functions for data filling, determining the class label and running the CF clusterizer
 template <>
 GPUdii() void GPUTPCNNClusterizerKernels::Thread<GPUTPCNNClusterizerKernels::runCfClusterizer>(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUSharedMemory& smem, processorType& processors, uint8_t sector, int8_t dtype, int8_t withMC, uint32_t batchStart)
@@ -476,7 +488,6 @@ GPUdii() void GPUTPCNNClusterizerKernels::Thread<GPUTPCNNClusterizerKernels::pub
     }
     return;
   }
-
   uint32_t rowIndex = 0;
   if (clusterOut != nullptr) {
     rowIndex = GPUTPCCFClusterizer::sortIntoBuckets(
@@ -492,6 +503,9 @@ GPUdii() void GPUTPCNNClusterizerKernels::Thread<GPUTPCNNClusterizerKernels::pub
   } else if (clusterer.mPclusterPosInRow) {
     rowIndex = clusterer.mPclusterPosInRow[full_glo_idx];
   }
+  if (clustererNN.mNnClusterizerUseMomentumVector && clusterer.mPclusterNNDirectionByRow != nullptr && rowIndex < clusterer.mNMaxClusterPerRow) {
+    clusterer.mPclusterNNDirectionByRow[peak.row() * clusterer.mNMaxClusterPerRow + rowIndex] = getClass1NNDirection(clustererNN, dtype, model_output_index);
+  }
   CPU_ONLY(labelAcc->commit(peak.row(), rowIndex, clusterer.mNMaxClusterPerRow));
 }
 
@@ -593,6 +607,7 @@ GPUdii() void GPUTPCNNClusterizerKernels::Thread<GPUTPCNNClusterizerKernels::pub
     }
     return;
   }
+  // setClass2NNDirection(clustererNN, dtype, model_output_index, myCluster);
 
   uint32_t rowIndex = 0;
   if (clusterOut != nullptr) {
@@ -646,6 +661,7 @@ GPUdii() void GPUTPCNNClusterizerKernels::Thread<GPUTPCNNClusterizerKernels::pub
     }
     return;
   }
+  // setClass2NNDirection(clustererNN, dtype, model_output_index, myCluster);
 
   if (clusterOut != nullptr) {
     rowIndex = GPUTPCCFClusterizer::sortIntoBuckets(
diff --git a/GPU/Workflow/src/GPUWorkflowSpec.cxx b/GPU/Workflow/src/GPUWorkflowSpec.cxx
index d928cfdd502c9..6b88d45207951 100644
--- a/GPU/Workflow/src/GPUWorkflowSpec.cxx
+++ b/GPU/Workflow/src/GPUWorkflowSpec.cxx
@@ -1259,7 +1259,6 @@ Inputs GPURecoWorkflowSpec::inputs()
     std::map<std::string, std::string> metadata;
     metadata["inputDType"] = nnClusterizerSettings.nnInferenceInputDType;              // FP16 or FP32
     metadata["outputDType"] = nnClusterizerSettings.nnInferenceOutputDType;            // FP16 or FP32
-    metadata["nnCCDBWithMomentum"] = nnClusterizerSettings.nnCCDBWithMomentum;         // 0, 1 -> Only for regression model
     metadata["nnCCDBLayerType"] = nnClusterizerSettings.nnCCDBClassificationLayerType; // FC, CNN
     metadata["nnCCDBInteractionRate"] = nnClusterizerSettings.nnCCDBInteractionRate;   // in kHz
     metadata["nnCCDBBeamType"] = nnClusterizerSettings.nnCCDBBeamType;                 // pp, pPb, PbPb
@@ -1286,6 +1285,7 @@ Inputs GPURecoWorkflowSpec::inputs()
       printSettings(metadata);
     }
 
+    metadata["nnCCDBWithMomentum"] = "0";
     if (mSpecConfig.nnEvalMode[0] == "c1") {
       metadata["nnCCDBEvalType"] = "classification_c1";
       convert_map_to_metadata(metadata, ccdb_metadata);
@@ -1297,6 +1297,7 @@ Inputs GPURecoWorkflowSpec::inputs()
       inputs.emplace_back("nn_classification_c2", gDataOriginTPC, "NNCLUSTERIZER_C2", 0, Lifetime::Condition, ccdbParamSpec(nnClusterizerSettings.nnCCDBPath + "/" + metadata["nnCCDBEvalType"], ccdb_metadata, 0));
     }
 
+    metadata["nnCCDBWithMomentum"] = nnClusterizerSettings.nnCCDBWithMomentum; // 0, 1 -> Only for regression model
     metadata["nnCCDBEvalType"] = "regression_c1";
     metadata["nnCCDBLayerType"] = nnClusterizerSettings.nnCCDBRegressionLayerType;
     convert_map_to_metadata(metadata, ccdb_metadata);