diff --git a/src/dfChemistryModel/DNNInferencer/DNNInferencer.cpp b/src/dfChemistryModel/DNNInferencer/DNNInferencer.cpp
index 0092bb6e..69aa784e 100644
--- a/src/dfChemistryModel/DNNInferencer/DNNInferencer.cpp
+++ b/src/dfChemistryModel/DNNInferencer/DNNInferencer.cpp
@@ -265,7 +265,7 @@ std::vector<std::vector<double>> DNNInferencer::Inference_multiDNNs(std::vector<
 
     std::chrono::steady_clock::time_point stop = std::chrono::steady_clock::now();
     std::chrono::duration<double> processingTime = std::chrono::duration_cast<std::chrono::duration<double>>(stop - start);
-    std::cout << "preInf time = " << processingTime.count() << std::endl;
+    // std::cout << "preInf time = " << processingTime.count() << std::endl;
     time_preInf += processingTime.count();
 
     // inference
@@ -288,7 +288,7 @@ std::vector<std::vector<double>> DNNInferencer::Inference_multiDNNs(std::vector<
 
     std::chrono::steady_clock::time_point stop1 = std::chrono::steady_clock::now();
     std::chrono::duration<double> processingTime1 = std::chrono::duration_cast<std::chrono::duration<double>>(stop1 - start1);
-    std::cout << "Inf time = " << processingTime1.count() << std::endl;
+    // std::cout << "Inf time = " << processingTime1.count() << std::endl;
     time_Inference += processingTime1.count();
 
     // generate outputTensor
@@ -322,7 +322,7 @@ std::vector<std::vector<double>> DNNInferencer::Inference_multiDNNs(std::vector<
 
     std::chrono::steady_clock::time_point stop3 = std::chrono::steady_clock::now();
     std::chrono::duration<double> processingTime3 = std::chrono::duration_cast<std::chrono::duration<double>>(stop3 - start3);
-    std::cout << "hot time = " << processingTime3.count() << std::endl;
+    // std::cout << "hot time = " << processingTime3.count() << std::endl;
     time_hot += processingTime3.count();
 
     std::vector<double> RRoutputs0(Youtputs0.data_ptr<double>(), Youtputs0.data_ptr<double>() + Youtputs0.numel());
@@ -333,13 +333,13 @@ std::vector<std::vector<double>> DNNInferencer::Inference_multiDNNs(std::vector<
 
     std::chrono::steady_clock::time_point stop2 = std::chrono::steady_clock::now();
     std::chrono::duration<double> processingTime2 = std::chrono::duration_cast<std::chrono::duration<double>>(stop2 - start2);
-    std::cout << "postInf time = " << processingTime2.count() << std::endl;
+    // std::cout << "postInf time = " << processingTime2.count() << std::endl;
     time_postInf += processingTime2.count();
 
-    std::cout << "preInf sum time = " << time_preInf << std::endl;
-    std::cout << "Inf sum time = " << time_Inference << std::endl;
-    std::cout << "postInf sum time = " << time_postInf << std::endl;
-    std::cout << "hot sum time = " << time_hot << std::endl;
+    // std::cout << "preInf sum time = " << time_preInf << std::endl;
+    // std::cout << "Inf sum time = " << time_Inference << std::endl;
+    // std::cout << "postInf sum time = " << time_postInf << std::endl;
+    // std::cout << "hot sum time = " << time_hot << std::endl;
 
     return results;
 }
diff --git a/src/dfChemistryModel/dfChemistryModel.C b/src/dfChemistryModel/dfChemistryModel.C
index 334be4d0..a0b4a7c9 100644
--- a/src/dfChemistryModel/dfChemistryModel.C
+++ b/src/dfChemistryModel/dfChemistryModel.C
@@ -99,7 +99,14 @@ Foam::dfChemistryModel<ThermoType>::dfChemistryModel
     )
 {
 #ifdef USE_LIBTORCH
+    useDNN = true;
+    if (!Qdot_.typeHeaderOk<volScalarField>())
+    {
+        useDNN = false;
+    }
+
     torchSwitch_ = this->lookupOrDefault("torch", false);
+    gpu_ = this->lookupOrDefault("GPU", false),
 
     torchModelName_ = this->lookupOrDefault("torchModel", word(""));
     torchModelName1_ = this->lookupOrDefault("torchModel1", word(""));
@@ -125,14 +132,23 @@ Foam::dfChemistryModel<ThermoType>::dfChemistryModel
     {
         if(!(Pstream::myProcNo() % coresPerGPU)) // Now is a master
         {
-            int CUDANo = (Pstream::myProcNo() / coresPerGPU) % GPUsPerNode;
-            std::string device_ = "cuda:" + std::to_string(CUDANo);
+            
             Info << "location 0" << endl;
             Info << "torchModelName1_ = " << torchModelName1_ << endl;
             torch::jit::script::Module torchModel1_ = torch::jit::load(torchModelName1_);
             torch::jit::script::Module torchModel2_ = torch::jit::load(torchModelName2_);
             torch::jit::script::Module torchModel3_ = torch::jit::load(torchModelName3_);
             Info << "location 1" << endl;
+            std::string device_;
+            if (gpu_)
+            {
+                int CUDANo = (Pstream::myProcNo() / coresPerGPU) % GPUsPerNode;
+                device_ = "cuda:" + std::to_string(CUDANo); 
+            }
+            else
+            {
+                device_ = "cpu";
+            }
             DNNInferencer DNNInferencer(torchModel1_, torchModel2_, torchModel3_, device_);
             DNNInferencer_ = DNNInferencer;
         }
@@ -140,6 +156,13 @@ Foam::dfChemistryModel<ThermoType>::dfChemistryModel
 #endif
 
 #ifdef USE_PYTORCH
+    useDNN = true;
+    if (!Qdot_.typeHeaderOk<volScalarField>())
+    {
+        useDNN = false;
+    }
+    gpu_ = this->lookupOrDefault("GPU", false),
+
     torchSwitch_ = this->lookupOrDefault("torch", false);
 
     Tact1_ = this->subDict("torchParameters1").lookupOrDefault("Tact", 700);
@@ -276,23 +299,39 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve
 #ifdef USE_LIBTORCH
     if(torchSwitch_)
     {
-        result = solve_DNN_GPU(deltaT);
+        if (useDNN)
+        {
+            result = solve_DNN(deltaT);
+        }
+        else
+        {
+            result = solve_CVODE(deltaT);
+            useDNN = true;
+        }
     }
     else
     {
-        result = solve_loadBalance(deltaT);
+        result = solve_CVODE(deltaT);
     }
 #elif USE_PYTORCH
     if(torchSwitch_)
     {
-        result = solve_DNN_GPU(deltaT);
+        if (useDNN)
+        {
+            result = solve_DNN(deltaT);
+        }
+        else
+        {
+            result = solve_CVODE(deltaT);
+            useDNN = true;
+        }
     }
     else
     {
-        result = solve_loadBalance(deltaT);
+        result = solve_CVODE(deltaT);
     }
 #else
-    result = solve_loadBalance(deltaT);
+    result = solve_CVODE(deltaT);
 #endif
 
     return result;
@@ -331,152 +370,6 @@ void Foam::dfChemistryModel<ThermoType>::setNumerics(Cantera::ReactorNet &sim)
 
 
 #ifdef USE_LIBTORCH
-template<class ThermoType>
-template<class DeltaTType>
-Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN
-(
-    const DeltaTType& deltaT
-)
-{
-    scalar deltaTMin = great;
-
-    if (!this->chemistry_)
-    {
-        return deltaTMin;
-    }
-
-    Info<<"=== begin solve_DNN === "<<endl;
-
-    // set variables
-    scalarList yPre_(mixture_.nSpecies());
-    scalarList yBCT_(mixture_.nSpecies());
-    scalarList u_(mixture_.nSpecies()+2); //plus T and p
-    Cantera::Reactor react;
-    double lambda = 0.1;
-
-    // load model
-    torch::DeviceType device_type;
-    torch::DeviceIndex index = -1;
-    if (0)
-    {
-        device_type = torch::kCUDA;
-        index = 0;
-    }
-    else
-    {
-        device_type = torch::kCPU;
-        index = 0;
-    }
-    torch::Device device(device_type, index);
-    if (!isFile(torchModelName_))
-    {
-        FatalErrorInFunction
-            << torchModelName_
-            << " doesn't exist!"
-            << exit(FatalError);
-    }
-    torch::jit::script::Module torchModel_ = torch::jit::load(torchModelName_, device);
-
-    std::vector<label> torch_cell;
-    label torch_cellname= 0;
-
-    // obtain the number of DNN cells
-    forAll(T_, cellI)
-    {
-        if ((Qdot_[cellI] >= Qdotact_) && (T_[cellI] >= Tact_))
-        {
-            torch_cell.push_back(cellI);
-        }
-    }
-
-    torch::Tensor inputs(torch::zeros({torch_cell.size(),mixture_.nSpecies()+2}));
-    forAll(T_, cellI)
-    {
-        scalar Ti = T_[cellI];
-        scalar pi = p_[cellI];
-        scalar rhoi = rho_[cellI];
-
-        if ((Qdot_[cellI] >= Qdotact_) && (Ti >= Tact_))
-        {
-            Qdot_[cellI] = 0.0;
-            // Info<<"Now is DNN "<<endl;
-
-            // set inputs
-            std::vector<double> inputs_;
-            inputs_.push_back((Ti - Xmu_[0])/Xstd_[0]);
-            inputs_.push_back((pi / 101325 - Xmu_[1])/Xstd_[1]);
-            for (size_t i=0; i<CanteraGas_->nSpecies(); ++i)
-            {
-                yPre_[i] = Y_[i][cellI];
-                yBCT_[i] = (pow(yPre_[i],lambda) - 1) / lambda; // function BCT
-            }
-            for (size_t i=0; i<CanteraGas_->nSpecies(); ++i)
-            {
-                inputs_.push_back((yBCT_[i] - Xmu_[i+2]) / Xstd_[i+2]);
-            }
-            inputs[torch_cellname] = torch::tensor(inputs_);
-            torch_cellname ++;
-            // Info << "cell_name = "<< cellI <<endl;
-        }
-        else
-        {
-            Qdot_[cellI] = 0.0;
-            for (size_t i=0; i<CanteraGas_->nSpecies(); ++i)
-            {
-                yPre_[i] = Y_[i][cellI];
-            }
-
-            CanteraGas_->setState_TPY(Ti, pi, yPre_.begin());
-            react.insert(mixture_.CanteraSolution());
-            react.setEnergy(0);
-
-            Cantera::ReactorNet sim;
-            sim.addReactor(react);
-            setNumerics(sim);
-            sim.advance(deltaT[cellI]);
-
-            CanteraGas_->getMassFractions(yTemp_.begin());
-
-            for (size_t i=0; i<CanteraGas_->nSpecies(); ++i)
-            {
-                RR_[i][cellI] = (yTemp_[i] - yPre_[i])*rhoi/deltaT[cellI];
-                Qdot_[cellI] -= hc_[i]*RR_[i][cellI];
-            }
-        }
-    }
-    // DNN
-    std::vector<torch::jit::IValue> INPUTS{inputs};
-    auto outputs_raw = torchModel_.forward(INPUTS);
-    auto outputs = outputs_raw.toTensor();
-
-    for(size_t cellI = 0; cellI<torch_cell.size(); cellI ++)
-    {
-        // update y
-        scalar Yt = 0;
-        for (size_t i=0; i<CanteraGas_->nSpecies(); ++i)
-        {
-            yPre_[i] = Y_[i][torch_cell[cellI]];
-            yTemp_[i] = Y_[i][torch_cell[cellI]];
-            yBCT_[i] = (pow(yPre_[i],lambda) - 1) / lambda; // function BCT
-        }
-        for (size_t i=0; i<(CanteraGas_->nSpecies()); ++i)//
-        {
-            u_[i+2] = outputs[cellI][i+2].item().to<double>()*Ystd_[i+2]+Ymu_[i+2];
-            yTemp_[i] = pow((yBCT_[i] + u_[i+2]*deltaT[cellI])*lambda+1,1/lambda);
-            Yt += yTemp_[i];
-        }
-        for (size_t i=0; i<CanteraGas_->nSpecies(); ++i)
-        {
-            yTemp_[i] = yTemp_[i] / Yt;
-            RR_[i][torch_cell[cellI]] = (yTemp_[i] - Y_[i][torch_cell[cellI]])*rho_[torch_cell[cellI]]/deltaT[cellI];
-            Qdot_[torch_cell[cellI]] -= hc_[i]*RR_[i][torch_cell[cellI]];
-        }
-    }
-
-    Info<<"=== end torch&ode-solve === "<<endl;
-    return deltaTMin;
-}
-
 template <class ThermoType>
 template<class DeltaTType>
 Foam::DynamicList<Foam::GpuProblem>
@@ -520,34 +413,34 @@ Foam::dfChemistryModel<ThermoType>::getGPUProblems
         // choose DNN module
         if ((Qdot_[cellI] < Qdotact2_) && (T_[cellI] < Tact2_) && ( T_[cellI] >= Tact1_))//choose1
         {
-            problem.DNNid = 0;
+            problem.DNNid = 0; 
             problemList.append(problem);
             continue;
         }
         if(((Qdot_[cellI] >= Qdotact2_) && (T_[cellI] < Tact2_)&&(T_[cellI] >= Tact1_))||((Qdot_[cellI] > Qdotact3_) && T_[cellI] > Tact2_))  //choose2
         {
-            problem.DNNid = 1;
+            problem.DNNid = 1; 
             problemList.append(problem);
             continue;
         }
         if  ((Qdot_[cellI] < Qdotact3_) && (T_[cellI] >= Tact2_) && (Qdot_[cellI]!=0)) //choose3
         {
-            problem.DNNid = 2;
+            problem.DNNid = 2; 
             problemList.append(problem);
             continue;
         }
 
     }
 
-    return problemList;
+    return problemList; 
 }
 
 template <class ThermoType>
 void Foam::dfChemistryModel<ThermoType>::getDNNinputs
 (
-    const Foam::DynamicBuffer<GpuProblem>& problemBuffer,
+    const Foam::DynamicBuffer<GpuProblem>& problemBuffer, 
     std::vector<label>& outputLength,
-    std::vector<std::vector<double>>& DNNinputs,
+    std::vector<std::vector<double>>& DNNinputs, 
     std::vector<Foam::DynamicBuffer<label>>& cellIDBuffer,
     std::vector<std::vector<label>>& problemCounter
 )
@@ -610,7 +503,7 @@ void Foam::dfChemistryModel<ThermoType>::getDNNinputs
                 counter2++;
                 cellIDList2.append(problemBuffer[i][cellI].cellid);
                 break;
-
+            
             default:
                 Info<<"invalid input"<<endl;
                 break;
@@ -638,7 +531,7 @@ void Foam::dfChemistryModel<ThermoType>::getDNNinputs
     // auto inputTensor1 = torch::tensor(inputsDNN1);
     // inputTensor1 = inputTensor1.reshape({length1, mixture_.nSpecies() + 3});
     // auto inputTensor2 = torch::tensor(inputsDNN2);
-    // inputTensor2 = inputTensor2.reshape({length2, mixture_.nSpecies() + 3});
+    // inputTensor2 = inputTensor2.reshape({length2, mixture_.nSpecies() + 3});    
 
     // set output
     outputLength = {length0, length1, length2};
@@ -646,11 +539,9 @@ void Foam::dfChemistryModel<ThermoType>::getDNNinputs
     cellIDBuffer = {cellIDList0Buffer, cellIDList1Buffer, cellIDList2Buffer};
     problemCounter = {problemCounter0, problemCounter1, problemCounter2};
 
-    std::cout<<"inputsDNN0 = "<<inputsDNN0.size()<<endl;
-    std::cout<<"inputsDNN1 = "<<inputsDNN1.size()<<endl;
-    std::cout<<"inputsDNN2 = "<<inputsDNN2.size()<<endl;
-
-    Info<<"get inputs successfully"<<endl;
+    std::cout<<"inputsDNN0 = "<<inputsDNN0.size()/10 << "\n";
+    std::cout<<"inputsDNN1 = "<<inputsDNN1.size()/10 << "\n";
+    std::cout<<"inputsDNN2 = "<<inputsDNN2.size()/10 << "\n";
 
     return;
 }
@@ -658,7 +549,7 @@ void Foam::dfChemistryModel<ThermoType>::getDNNinputs
 template <class ThermoType>
 void Foam::dfChemistryModel<ThermoType>::updateSolutionBuffer
 (
-    Foam::DynamicBuffer<Foam::GpuSolution>& solutionBuffer,
+    Foam::DynamicBuffer<Foam::GpuSolution>& solutionBuffer, 
     const std::vector<std::vector<double>>& results,
     const std::vector<Foam::label>& outputLength,
     const std::vector<Foam::DynamicBuffer<Foam::label>>& cellIDBuffer,
@@ -712,7 +603,7 @@ void Foam::dfChemistryModel<ThermoType>::updateSolutionBuffer
 
 template <class ThermoType>
 template <class DeltaTType>
-Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
+Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN(
     const DeltaTType &deltaT)
 {
     scalar deltaTMin = great;
@@ -723,7 +614,15 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
     }
 
     std::chrono::steady_clock::time_point start = std::chrono::steady_clock::now();
-    Info << "=== begin solve_DNN_CUDA === " << endl;
+    Info << "=== begin solve_DNN === " << endl;
+    if (gpu_)
+    {
+        Info << "now DNN inference is conducted on GPU" << endl;
+    }
+    else
+    {
+        Info << "now DNN inference is conducted on CPU" << endl;
+    }
 
     /*=============================gather problems=============================*/
     DynamicList<GpuProblem> problemList = getGPUProblems(deltaT);
@@ -731,19 +630,19 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
     /*==============================send problems==============================*/
     std::chrono::steady_clock::time_point start2 = std::chrono::steady_clock::now();
 
-    PstreamBuffers pBufs(Pstream::commsTypes::nonBlocking);
-    if (Pstream::myProcNo() % coresPerGPU) //for slave
+    PstreamBuffers pBufs(Pstream::commsTypes::nonBlocking); 
+    if (Pstream::myProcNo() % coresPerGPU) //for slave 
     {
         UOPstream send((Pstream::myProcNo()/coresPerGPU)*coresPerGPU, pBufs);// sending problem to master
         send << problemList;
     }
     pBufs.finishedSends();
-
+    
     DynamicBuffer<GpuSolution> solutionBuffer;
 
     std::chrono::steady_clock::time_point stop2 = std::chrono::steady_clock::now();
     std::chrono::duration<double> processingTime2 = std::chrono::duration_cast<std::chrono::duration<double>>(stop2 - start2);
-    std::cout << "sendProblemTime = " << processingTime2.count() << std::endl;
+    // std::cout << "sendProblemTime = " << processingTime2.count() << std::endl;
     time_sendProblem_ += processingTime2.count();
 
     /*=============================submaster work start=============================*/
@@ -751,7 +650,7 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
     {
         std::chrono::steady_clock::time_point start1 = std::chrono::steady_clock::now();
         std::chrono::steady_clock::time_point start3 = std::chrono::steady_clock::now();
-
+        
         label problemSize = 0; // problemSize is defined to debug
         DynamicBuffer<GpuProblem> problemBuffer(coresPerGPU);//each submaster init a local problemBuffer TODO:rename it
 
@@ -769,7 +668,7 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
 
         std::chrono::steady_clock::time_point stop3 = std::chrono::steady_clock::now();
         std::chrono::duration<double> processingTime3 = std::chrono::duration_cast<std::chrono::duration<double>>(stop3 - start3);
-        std::cout << "RecvProblemTime = " << processingTime3.count() << std::endl;
+        // std::cout << "RecvProblemTime = " << processingTime3.count() << std::endl;
         time_RecvProblem_ += processingTime3.count();
 
         /*==============================construct DNN inputs==============================*/
@@ -782,7 +681,7 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
         getDNNinputs(problemBuffer, outputLength, DNNinputs, cellIDBuffer, problemCounter);
         std::chrono::steady_clock::time_point stop5 = std::chrono::steady_clock::now();
         std::chrono::duration<double> processingTime5 = std::chrono::duration_cast<std::chrono::duration<double>>(stop5 - start5);
-        std::cout << "getDNNinputsTime = " << processingTime5.count() << std::endl;
+        // std::cout << "getDNNinputsTime = " << processingTime5.count() << std::endl;
         time_getDNNinputs_ += processingTime5.count();
 
         /*=============================inference via DNNInferencer=============================*/
@@ -792,7 +691,7 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
 
         std::chrono::steady_clock::time_point stop7 = std::chrono::steady_clock::now();
         std::chrono::duration<double> processingTime7 = std::chrono::duration_cast<std::chrono::duration<double>>(stop7 - start7);
-        std::cout << "DNNinferenceTime = " << processingTime7.count() << std::endl;
+        // std::cout << "DNNinferenceTime = " << processingTime7.count() << std::endl;
         time_DNNinference_ += processingTime7.count();
 
         /*=============================construct solutions=============================*/
@@ -802,12 +701,12 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
 
         std::chrono::steady_clock::time_point stop6 = std::chrono::steady_clock::now();
         std::chrono::duration<double> processingTime6 = std::chrono::duration_cast<std::chrono::duration<double>>(stop6 - start6);
-        std::cout << "updateSolutionBufferTime = " << processingTime6.count() << std::endl;
+        // std::cout << "updateSolutionBufferTime = " << processingTime6.count() << std::endl;
         time_updateSolutionBuffer_ += processingTime6.count();
-
+        
         std::chrono::steady_clock::time_point stop1 = std::chrono::steady_clock::now();
         std::chrono::duration<double> processingTime1 = std::chrono::duration_cast<std::chrono::duration<double>>(stop1 - start1);
-        std::cout << "submasterTime = " << processingTime1.count() << std::endl;
+        // std::cout << "submasterTime = " << processingTime1.count() << std::endl;  
         time_submaster_ += processingTime1.count();
     }
 
@@ -821,22 +720,22 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
         finalList = solutionBuffer[0];
         for (label i = 1; i < coresPerGPU; i++)
         {
-
+            
             UOPstream send(i + Pstream::myProcNo(), pBufs2);
             send << solutionBuffer[i];
         }
     }
     pBufs2.finishedSends();
-
+    
     if (Pstream::myProcNo() % coresPerGPU)
     {
         UIPstream recv((Pstream::myProcNo()/coresPerGPU)*coresPerGPU, pBufs2);
         recv >> finalList;
     }
-
+   
     std::chrono::steady_clock::time_point stop4 = std::chrono::steady_clock::now();
     std::chrono::duration<double> processingTime4 = std::chrono::duration_cast<std::chrono::duration<double>>(stop4 - start4);
-    std::cout << "SendRecvSolutionTime = " << processingTime4.count() << std::endl;
+    // std::cout << "SendRecvSolutionTime = " << processingTime4.count() << std::endl;
     time_sendRecvSolution_ += processingTime4.count();
 
     /*=============================update RR fields=============================*/
@@ -850,11 +749,11 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
         }
     }
 
-    Info << "=== end torch&ode-CUDAsolve === " << endl;
+    Info << "=== end solve_DNN === " << endl;
 
     std::chrono::steady_clock::time_point stop = std::chrono::steady_clock::now();
     std::chrono::duration<double> processingTime = std::chrono::duration_cast<std::chrono::duration<double>>(stop - start);
-    std::cout << "allSolveTime = " << processingTime.count() << std::endl;
+    // std::cout << "allSolveTime = " << processingTime.count() << std::endl;
     time_allsolve_ += processingTime.count();
 
     return deltaTMin;
@@ -1078,7 +977,7 @@ Foam::dfChemistryModel<ThermoType>::getProblems
     forAll(T, celli)
     {
         {
-            for(label i = 0; i < (int)(mixture_.nSpecies()); i++)
+            for(label i = 0; i < mixture_.nSpecies(); i++)
             {
                 yTemp_[i] = Y_[i][celli];
             }
@@ -1155,7 +1054,7 @@ Foam::dfChemistryModel<ThermoType>::updateReactionRates
         for(const auto& solution : array)
         {
 
-            for(label j = 0; j < int(mixture_.nSpecies()); j++)
+            for(label j = 0; j < mixture_.nSpecies(); j++)
             {
                 this->RR_[j][solution.cellid] = solution.RRi[j];
             }
@@ -1193,7 +1092,7 @@ Foam::dfChemistryModel<ThermoType>::createBalancer()
 
 template <class ThermoType>
 template <class DeltaTType>
-Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_loadBalance
+Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_CVODE
 (
     const DeltaTType& deltaT
 )
@@ -1452,23 +1351,23 @@ Foam::dfChemistryModel<ThermoType>::getGPUProblems
         // choose DNN module
         if ((Qdot_[cellI] < Qdotact2_) && (T_[cellI] < Tact2_) && ( T_[cellI] >= Tact1_))//choose1
         {
-            problem.DNNid = 0;
+            problem.DNNid = 0; 
             problemList.append(problem);
             continue;
         }
         if(((Qdot_[cellI] >= Qdotact2_) && (T_[cellI] < Tact2_)&&(T_[cellI] >= Tact1_))||((Qdot_[cellI] > Qdotact3_) && T_[cellI] > Tact2_))  //choose2
         {
-            problem.DNNid = 1;
+            problem.DNNid = 1; 
             problemList.append(problem);
             continue;
         }
         if  ((Qdot_[cellI] < Qdotact3_) && (T_[cellI] >= Tact2_) && (Qdot_[cellI]!=0)) //choose3
         {
-            problem.DNNid = 2;
+            problem.DNNid = 2; 
             problemList.append(problem);
             continue;
         }
-
+    
     }
 
     return problemList;
@@ -1570,11 +1469,9 @@ void Foam::dfChemistryModel<ThermoType>::getDNNinputs
     cellIDBuffer = {cellIDList0Buffer, cellIDList1Buffer, cellIDList2Buffer};
     problemCounter = {problemCounter0, problemCounter1, problemCounter2};
 
-    std::cout << "inputsDNN0 = " << inputsDNN0.size();
-    std::cout << "inputsDNN1 = " << inputsDNN1.size();
-    std::cout << "inputsDNN2 = " << inputsDNN2.size();
-
-    Info<<"get inputs successfully"<<endl;
+    std::cout<<"inputsDNN0 = "<<inputsDNN0.size()/10 << "\n";
+    std::cout<<"inputsDNN1 = "<<inputsDNN1.size()/10 << "\n";
+    std::cout<<"inputsDNN2 = "<<inputsDNN2.size()/10 << "\n";
 
     return;
 }
@@ -1640,7 +1537,7 @@ void Foam::dfChemistryModel<ThermoType>::updateSolutionBuffer
 
 template <class ThermoType>
 template <class DeltaTType>
-Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
+Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN(
     const DeltaTType &deltaT)
 {
     scalar deltaTMin = great;
@@ -1651,7 +1548,16 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
     }
 
     std::chrono::steady_clock::time_point start = std::chrono::steady_clock::now();
-    Info << "=== begin solve_DNN_GPU === " << endl;
+    Info << "=== begin solve_DNN === " << endl;
+    if (gpu_)
+    {
+        Info << "now DNN inference is conducted on GPU" << endl;
+    }
+    else
+    {
+        Info << "now DNN inference is conducted on CPU" << endl;
+    }
+    
 
     /*=============================gather problems=============================*/
     DynamicList<GpuProblem> problemList = getGPUProblems(deltaT);
@@ -1671,7 +1577,7 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
 
     std::chrono::steady_clock::time_point stop2 = std::chrono::steady_clock::now();
     std::chrono::duration<double> processingTime2 = std::chrono::duration_cast<std::chrono::duration<double>>(stop2 - start2);
-    std::cout << "sendProblemTime = " << processingTime2.count() << std::endl;
+    // std::cout << "sendProblemTime = " << processingTime2.count() << std::endl;
     time_sendProblem_ += processingTime2.count();
 
     /*=============================submaster work start=============================*/
@@ -1697,7 +1603,7 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
 
         std::chrono::steady_clock::time_point stop3 = std::chrono::steady_clock::now();
         std::chrono::duration<double> processingTime3 = std::chrono::duration_cast<std::chrono::duration<double>>(stop3 - start3);
-        std::cout << "RecvProblemTime = " << processingTime3.count() << std::endl;
+        // std::cout << "RecvProblemTime = " << processingTime3.count() << std::endl;
         time_RecvProblem_ += processingTime3.count();
 
         /*==============================construct DNN inputs==============================*/
@@ -1710,7 +1616,7 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
         getDNNinputs(problemBuffer, outputLength, DNNinputs, cellIDBuffer, problemCounter);
         std::chrono::steady_clock::time_point stop5 = std::chrono::steady_clock::now();
         std::chrono::duration<double> processingTime5 = std::chrono::duration_cast<std::chrono::duration<double>>(stop5 - start5);
-        std::cout << "getDNNinputsTime = " << processingTime5.count() << std::endl;
+        // std::cout << "getDNNinputsTime = " << processingTime5.count() << std::endl;
         time_getDNNinputs_ += processingTime5.count();
 
         /*=============================inference via pybind11=============================*/
@@ -1723,7 +1629,7 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
 
         std::chrono::steady_clock::time_point stop8 = std::chrono::steady_clock::now();
         std::chrono::duration<double> processingTime8 = std::chrono::duration_cast<std::chrono::duration<double>>(stop8 - start8);
-        std::cout << "vec2ndarrayTime = " << processingTime8.count() << std::endl;
+        // std::cout << "vec2ndarrayTime = " << processingTime8.count() << std::endl;
         time_vec2ndarray_ += processingTime8.count();
 
         pybind11::module_ call_torch = pybind11::module_::import("inference"); // import python file
@@ -1735,12 +1641,12 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
 
         std::chrono::steady_clock::time_point stop9 = std::chrono::steady_clock::now();
         std::chrono::duration<double> processingTime9 = std::chrono::duration_cast<std::chrono::duration<double>>(stop9 - start9);
-        std::cout << "pythonTime = " << processingTime9.count() << std::endl;
+        // std::cout << "pythonTime = " << processingTime9.count() << std::endl;
         time_python_ += processingTime9.count();
 
         std::chrono::steady_clock::time_point stop7 = std::chrono::steady_clock::now();
         std::chrono::duration<double> processingTime7 = std::chrono::duration_cast<std::chrono::duration<double>>(stop7 - start7);
-        std::cout << "DNNinferenceTime = " << processingTime7.count() << std::endl;
+        // std::cout << "DNNinferenceTime = " << processingTime7.count() << std::endl;
         time_DNNinference_ += processingTime7.count();
 
         /*=============================construct solutions=============================*/
@@ -1748,12 +1654,12 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
         updateSolutionBuffer(solutionBuffer, star, outputLength, cellIDBuffer, problemCounter);
         std::chrono::steady_clock::time_point stop6 = std::chrono::steady_clock::now();
         std::chrono::duration<double> processingTime6 = std::chrono::duration_cast<std::chrono::duration<double>>(stop6 - start6);
-        std::cout << "updateSolutionBufferTime = " << processingTime6.count() << std::endl;
+        // std::cout << "updateSolutionBufferTime = " << processingTime6.count() << std::endl;
         time_updateSolutionBuffer_ += processingTime6.count();
 
         std::chrono::steady_clock::time_point stop1 = std::chrono::steady_clock::now();
         std::chrono::duration<double> processingTime1 = std::chrono::duration_cast<std::chrono::duration<double>>(stop1 - start1);
-        std::cout << "submasterTime = " << processingTime1.count() << std::endl;
+        // std::cout << "submasterTime = " << processingTime1.count() << std::endl;
         time_submaster_ += processingTime1.count();
     }
 
@@ -1780,7 +1686,7 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
 
     std::chrono::steady_clock::time_point stop4 = std::chrono::steady_clock::now();
     std::chrono::duration<double> processingTime4 = std::chrono::duration_cast<std::chrono::duration<double>>(stop4 - start4);
-    std::cout << "SendRecvSolutionTime = " << processingTime4.count() << std::endl;
+    // std::cout << "SendRecvSolutionTime = " << processingTime4.count() << std::endl;
     time_sendRecvSolution_ += processingTime4.count();
 
     /*=============================update RR fields=============================*/
@@ -1794,11 +1700,11 @@ Foam::scalar Foam::dfChemistryModel<ThermoType>::solve_DNN_GPU(
         }
     }
 
-    Info << "=== end torch&ode-CUDAsolve === " << endl;
+    Info << "=== end solve_DNN === " << endl;
 
     std::chrono::steady_clock::time_point stop = std::chrono::steady_clock::now();
     std::chrono::duration<double> processingTime = std::chrono::duration_cast<std::chrono::duration<double>>(stop - start);
-    std::cout << "allSolveTime = " << processingTime.count() << std::endl;
+    // std::cout << "allSolveTime = " << processingTime.count() << std::endl;
     time_allsolve_ += processingTime.count();
 
     return deltaTMin;
diff --git a/src/dfChemistryModel/dfChemistryModel.H b/src/dfChemistryModel/dfChemistryModel.H
index 369f107e..bb94bb38 100644
--- a/src/dfChemistryModel/dfChemistryModel.H
+++ b/src/dfChemistryModel/dfChemistryModel.H
@@ -128,7 +128,10 @@ public IOdictionary
 
 
 #ifdef USE_LIBTORCH
+    bool useDNN;
+
     Switch torchSwitch_;
+    Switch gpu_;
     DNNInferencer DNNInferencer_;
     int coresPerGPU;
     int GPUsPerNode;
@@ -165,7 +168,10 @@ public IOdictionary
 #endif
 
 #ifdef USE_PYTORCH
+    bool useDNN;
+
     Switch torchSwitch_;
+    Switch gpu_;
     scalar Tact1_;
     scalar Qdotact1_;
     scalar Tact2_;
@@ -212,9 +218,6 @@ public IOdictionary
         scalar canteraSolve(const DeltaTType& deltaT);
 
 #ifdef USE_LIBTORCH
-        //- DNN solver
-        template<class DeltaTType>
-        scalar solve_DNN(const DeltaTType& deltaT);
 
         //- DNN solver that support the use of GPU (single card & single CPU core)
         template<class DeltaTType>
@@ -222,9 +225,9 @@ public IOdictionary
 
         //- DNN solver that support the use of GPU (multi-cards)
         template<class DeltaTType>
-        scalar solve_DNN_GPU(const DeltaTType& deltaT);
+        scalar solve_DNN(const DeltaTType& deltaT);
 
-        /*============= private methods in solve_CVODE =============*/
+        /*============= private methods in solve_DNN =============*/
         //- get a list of GPU problems to be solved
         template<class DeltaTType>
         DynamicList<GpuProblem> getGPUProblems(const DeltaTType& deltaT);
@@ -245,9 +248,9 @@ public IOdictionary
         scalar solve_DNN_GPU_oneCore(const DeltaTType& deltaT);
 
         template<class DeltaTType>
-        scalar solve_DNN_GPU(const DeltaTType& deltaT);
+        scalar solve_DNN(const DeltaTType& deltaT);
 
-        /*=======================private methods for solve_DNN_GPU=======================*/
+        /*=======================private methods for solve_DNN=======================*/
         //- get a list of GPU problems to be solved
         template<class DeltaTType>
         DynamicList<GpuProblem> getGPUProblems(const DeltaTType& deltaT);
@@ -286,7 +289,7 @@ public IOdictionary
 
         //- Solve the reaction system with DLB algorithm
         template<class DeltaTType>
-        scalar solve_loadBalance(const DeltaTType& deltaT);
+        scalar solve_CVODE(const DeltaTType& deltaT);
 
         //- Output logFiles
         Foam::autoPtr<Foam::OFstream> logFile(const word& name) const
@@ -361,6 +364,16 @@ public:
     double time_vec2ndarray() {return time_vec2ndarray_;}
     double time_python() {return time_python_;}
 #endif
+#ifdef USE_LIBTORCH
+    double time_allsolve() {return time_allsolve_;}
+    double time_submaster() {return time_submaster_;}
+    double time_sendProblem() {return time_sendProblem_;}
+    double time_RecvProblem() {return time_RecvProblem_;}
+    double time_sendRecvSolution() {return time_sendRecvSolution_;}
+    double time_getDNNinputs() {return time_getDNNinputs_;}
+    double time_DNNinference() {return time_DNNinference_;}
+    double time_updateSolutionBuffer() {return time_updateSolutionBuffer_;}
+#endif
 
 };