diff --git a/src/dfChemistryModel/pytorchFunctions.H b/src/dfChemistryModel/pytorchFunctions.H
index 535a7418..af2c224d 100644
--- a/src/dfChemistryModel/pytorchFunctions.H
+++ b/src/dfChemistryModel/pytorchFunctions.H
@@ -275,9 +275,9 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN(const DeltaTType &del
         pybind11::module_ call_torch = pybind11::module_::import("inference"); // import python file
         pybind11::object result = call_torch.attr("inference")(vec0, vec1, vec2); // call python function
         const double* star = result.cast<pybind11::array_t<double>>().data();
-        std::vector<double> outputsVec0(star, star+outputLength[0] * 7); //the float number is sample_length*sample_number
-        std::vector<double> outputsVec1(star+outputLength[0] * 7, star+outputLength[1] * 7);
-        std::vector<double> outputsVec2(star+outputLength[1] * 7, star+outputLength[2] * 7);
+        std::vector<double> outputsVec0(star, star+outputLength[0] * mixture_.nSpecies()); //the float number is sample_length*sample_number
+        std::vector<double> outputsVec1(star+outputLength[0] * mixture_.nSpecies(), star+outputLength[1] * mixture_.nSpecies());
+        std::vector<double> outputsVec2(star+outputLength[1] * mixture_.nSpecies(), star+outputLength[2] * mixture_.nSpecies());
         std::vector<std::vector<double>> results = {outputsVec0, outputsVec1, outputsVec2};
         updateSolutionBuffer(solutionBuffer, results, cellIDBuffer, problemCounter);
         DynamicList<GpuSolution> finalList;
@@ -301,4 +301,4 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN(const DeltaTType &del
     time_allsolve_ += processingTime.count();
 
     return deltaTMin;
-}
\ No newline at end of file
+}
diff --git a/src/dfChemistryModel/torchFunctions.H b/src/dfChemistryModel/torchFunctions.H
index 1639d287..f7d9a83c 100644
--- a/src/dfChemistryModel/torchFunctions.H
+++ b/src/dfChemistryModel/torchFunctions.H
@@ -149,9 +149,9 @@ void Foam::dfChemistryModel<ThermoType>::getDNNinputs
 
     if (gpulog_)
     {
-        std::cout<<"inputsDNN0 = "<<inputsDNN0.size()/10 << "\n";
-        std::cout<<"inputsDNN1 = "<<inputsDNN1.size()/10 << "\n";
-        std::cout<<"inputsDNN2 = "<<inputsDNN2.size()/10 << "\n";
+        std::cout<<"inputsDNN0 = "<<inputsDNN0.size()/(mixture_.nSpecies()+3) << "\n";
+        std::cout<<"inputsDNN1 = "<<inputsDNN1.size()/(mixture_.nSpecies()+3) << "\n";
+        std::cout<<"inputsDNN2 = "<<inputsDNN2.size()/(mixture_.nSpecies()+3) << "\n";
     }
 
     return;