diff --git a/docs/src/manual/pino_ode.md b/docs/src/manual/pino_ode.md
new file mode 100644
index 000000000..379f84c46
--- /dev/null
+++ b/docs/src/manual/pino_ode.md
@@ -0,0 +1,21 @@
+# Physics-Informed Neural operator for solve ODEs
+
+```@docs
+PINOODE
+```
+
+```@docs
+TRAINSET
+```
+
+```@docs
+PINOsolution
+```
+
+```@docs
+OperatorLearning
+```
+
+```@docs
+EquationSolving
+```
diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md
index 721b63c34..2e0ab0ed4 100644
--- a/docs/src/tutorials/pino_ode.md
+++ b/docs/src/tutorials/pino_ode.md
@@ -2,21 +2,24 @@
 
 This tutorial is an introduction to using physics-informed neural operator (PINOs) for solving family of parametric ordinary diferential equations (ODEs). 
 
+#TODO two phase
 
-## Solving a family of parametric ODE.
+## Operator Learning  for a family of parametric ODE.
 
 ```@example pino
 using Test
 using OrdinaryDiffEq, OptimizationOptimisers
 using Lux
 using Statistics, Random
-using NeuralOperators
+# using NeuralOperators
 using NeuralPDE
 
 linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p)
 linear = (u, p, t) -> cos(p * t)
 tspan = (0.0f0, 2.0f0)
 u0 = 0.0f0
+p = pi / 2f0
+prob = ODEProblem(linear, u0, tspan, p)
 ```
 
 Generate a dataset for learning a given family of ODEs where the parameter 'a' is varied. The dataset is generated by solving the ODE for different values of 'a' and storing the solutions. The dataset is then used to train the PINO model:
@@ -34,24 +37,33 @@ as = [Float32(i) for i in range(0.1, stop = pi / 2, length = batch_size)]
 u_output_ = zeros(Float32, 1, instances_size, batch_size)
 prob_set = []
 for (i, a_i) in enumerate(as)
-    prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i)
-    sol1 = solve(prob, Tsit5(); saveat = 0.0204)
+    prob_ = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i)
+    sol1 = solve(prob_, Tsit5(); saveat = 0.0204)
     reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1))
-    push!(prob_set, prob)
+    push!(prob_set, prob_)
     u_output_[:, :, i] = reshape_sol
 end
-train_set = TRAINSET(prob_set, u_output_);
+train_set = TRAINSET(prob_set, u_output_)
 ```
 
-Here it used the PINO method to train the given family of parametric ODEs. 
+Here it used the PINO method to learning operator of the given family of parametric ODEs. 
 
 ```@example pino
-prob = ODEProblem(linear, u0, tspan, 0)
-flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,),
-    σ = gelu)
-opt = OptimizationOptimisers.Adam(0.03)
-alg = PINOODE(flat_no, opt, train_set; is_data_loss = true, is_physics_loss = true)
-pino_solution = solve(prob, alg, verbose = false, maxiters = 1000)
+chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ),
+    Lux.Dense(16, 16, Lux.σ),
+    Lux.Dense(16, 16, Lux.σ),
+    Lux.Dense(16, 32, Lux.σ),
+    Lux.Dense(32, 32, Lux.σ),
+    Lux.Dense(32, 1))
+# flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,),
+#     σ = gelu)
+
+opt = OptimizationOptimisers.Adam(0.01)
+pino_phase = OperatorLearning(train_set, is_data_loss = true, is_physics_loss = true)
+
+alg = PINOODE(chain, opt, pino_phase)
+pino_solution = solve(
+    prob, alg, verbose = false, maxiters = 3000)
 predict = pino_solution.predict
 ground = u_output_
 ```
@@ -59,7 +71,35 @@ ground = u_output_
 Now let's compare the predictions from the learned operator with the ground truth solution which is obtained early by numerically solving the parametric ODE. Where 'i' is the index of the parameter 'a' in the dataset.
 
 ```@example pino
-i=1
+using Plots
+i=45
 plot(predict[1, :, i], label = "Predicted")
 plot!(ground[1, :, i], label = "Ground truth")
+```
+
+Now to move on the stage of solving a certain equation using a trained operator and physics
+
+## Solve ODE using learned operator family of parametric ODE for fine tuning.
+```@example pino
+dt = (t_end - t0) / instances_size
+pino_phase = EquationSolving(dt, pino_solution)
+chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ),
+    Lux.Dense(16, 16, Lux.σ),
+    Lux.Dense(16, 32, Lux.σ),
+    Lux.Dense(32, 1))
+alg = PINOODE(chain, opt, pino_phase)
+fine_tune_solution = solve( prob, alg, verbose = false, maxiters = 2000)
+
+fine_tune_predict = fine_tune_solution.predict
+operator_predict = pino_solution.phi(
+    fine_tune_solution.input_data_set, pino_solution.res.u)
+ground_fine_tune = linear_analytic.(u0, p, fine_tune_solution.input_data_set[[1], :, :])
+```
+
+Compare prediction with ground truth.
+
+```@example pino
+plot(operator_predict[1, :, 1], label = "operator_predict")
+plot!(fine_tune_predict[1, :, 1], label = "fine_tune_predict")
+plot!(ground_fine_tune[1, :, 1], label = "Ground truth")
 ```
\ No newline at end of file
diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl
index c1e00a2d4..15fcd4eb6 100644
--- a/src/pino_ode_solve.jl
+++ b/src/pino_ode_solve.jl
@@ -1,11 +1,12 @@
+"""
+TRAINSET(input_data, output_data; isu0 = false)
+## Positional Arguments
+* input_data: variables set 'a' of equation (for example initial condition {u(t0 x)} or parameter p
+* output_data: set of solutions u(t){a} corresponding parameter 'a'.
+"""
 struct TRAINSET{}
     input_data::Vector{ODEProblem}
     output_data::Array
-    isu0::Bool
-end
-
-function TRAINSET(input_data, output_data; isu0 = false)
-    TRAINSET(input_data, output_data, isu0)
 end
 
 mutable struct PINOPhi{C, T, U, S}
@@ -18,6 +19,14 @@ mutable struct PINOPhi{C, T, U, S}
     end
 end
 
+"""
+PINOsolution(predict, res, phi, input_data_set)
+## Positional Arguments
+* predict: The predicted solution.
+* res: The optimization solution.
+* phi: The solution operator.
+* input_data_set: The input data set.
+"""
 struct PINOsolution{}
     predict::Array
     res::SciMLBase.OptimizationSolution
@@ -26,25 +35,48 @@ struct PINOsolution{}
 end
 
 abstract type PINOPhases end
+
+"""
+OperatorLearning(train_set; is_data_loss = true, is_physics_loss = true)
+## Positional Arguments
+* train_set: Contains 'input data' - set of parameters 'a' and output data - set of solutions
+* u(t){a} corresponding initial conditions 'u0'.
+
+## Keyword Arguments
+* is_data_loss: Includes or off a loss function for training on the data set.
+* is_physics_loss: Includes or off loss function training on physics-informed approach.
+"""
 struct OperatorLearning <: PINOPhases
+    train_set::TRAINSET
     is_data_loss::Bool
     is_physics_loss::Bool
 end
-function OperatorLearning(; is_data_loss = true, is_physics_loss = true)
-    OperatorLearning(is_data_loss, is_physics_loss)
+function OperatorLearning(train_set; is_data_loss = true, is_physics_loss = true)
+    OperatorLearning(train_set, is_data_loss, is_physics_loss)
 end
+
+"""
+EquationSolving(dt, pino_solution)
+## Positional Arguments
+* dt: The time step.
+* pino_solution: Contains the solution of the operator learning phase.
+"""
 struct EquationSolving <: PINOPhases
+    dt::Number
     pino_solution::PINOsolution
+    is_finetune_loss::Bool
+    is_physics_loss::Bool
+end
+
+function EquationSolving(dt, pino_solution; is_finetune_loss = true, is_physics_loss = true)
+    EquationSolving(dt, pino_solution, is_finetune_loss, is_physics_loss)
 end
 
 """
    PINOODE(chain,
     OptimizationOptimisers.Adam(0.1),
-    train_set,
-    is_data_loss =true,
-    is_physics_loss =true,
+    pino_phase;
     init_params,
-    #TODO update docstring
     kwargs...)
 
 The method is that combine training data and physics constraints
@@ -54,14 +86,12 @@ to learn the solution operator of a given family of parametric Ordinary Differen
 * `chain`: A neural network architecture, defined as a `Lux.AbstractExplicitLayer` or `Flux.Chain`.
           `Flux.Chain` will be converted to `Lux` using `Lux.transform`.
 * `opt`: The optimizer to train the neural network.
-* `train_set`: Contains 'input data' - sr of parameters 'a' and output data - set of solutions
- u(t){a} corresponding initial conditions 'u0'.
+* `pino_phase`: The phase of the PINN algorithm, either `OperatorLearning` or `EquationSolving`.
 
 ## Keyword Arguments
-* `is_data_loss` Includes or off a loss function for training on the data set.
-* `is_physics_loss`: Includes or off loss function training on physics-informed approach.
 * `init_params`: The initial parameter of the neural network. By default, this is `nothing`
   which thus uses the random initialization provided by the neural network library.
+* isu0: If true, the input data set contains initial conditions 'u0'.
 * `kwargs`: Extra keyword arguments are splatted to the Optimization.jl `solve` call.
 
 ## References
@@ -70,21 +100,21 @@ Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Eq
 struct PINOODE{C, O, P, K} <: DiffEqBase.AbstractODEAlgorithm
     chain::C
     opt::O
-    train_set::TRAINSET
     pino_phase::PINOPhases
     init_params::P
+    isu0::Bool
     kwargs::K
 end
 
 function PINOODE(chain,
         opt,
-        train_set,
         pino_phase;
         init_params = nothing,
+        isu0 = false,
         kwargs...)
     #TODO fnn transform check
     !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain))
-    PINOODE(chain, opt, train_set, pino_phase, init_params, kwargs)
+    PINOODE(chain, opt, pino_phase, init_params, isu0, kwargs)
 end
 
 function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer,
@@ -132,23 +162,25 @@ function l₂loss(𝐲̂, 𝐲)
 end
 
 function physics_loss(phi::PINOPhi{C, T, U},
+        prob::ODEProblem,
         θ,
         ts::AbstractArray,
-        train_set::TRAINSET,
-        input_data_set) where {C, T, U}
-    prob_set, _ = train_set.input_data, train_set.output_data
-    f = prob_set[1].f
-    p = prob_set[1].p
+        prob_set::Array,
+        input_data_set,
+        isu0:: Bool) where {C, T, U}
+    f = prob.f
     out_ = phi(input_data_set, θ)
     ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), ts)
-    if train_set.isu0 == true
+    if isu0 == true
+        p = prob.p
         fs = f.f.(out_, p, ts)
     else
         ps = [prob.p for prob in prob_set]
-        if p isa Number
+        first_p = ps[1]
+        if first_p isa Number
             fs = cat(
                 [f.f.(out_[:, :, [i]], p, ts) for (i, p) in enumerate(ps)]..., dims = 3)
-        elseif p isa Vector
+        elseif first_p isa Vector
             fs = cat(
                 [reduce(
                      hcat, [f.f(out_[:, j, [i]], p, ts) for j in axes(out_[:, :, [i]], 2)])
@@ -163,14 +195,13 @@ end
 
 function data_loss(phi::PINOPhi{C, T, U},
         θ,
-        train_set::TRAINSET,
+        output_data::Array,
         input_data_set) where {C, T, U}
-    _, output_data = train_set.input_data, train_set.output_data
     output_data = adapt(parameterless_type(ComponentArrays.getdata(θ)), output_data)
     l₂loss(phi(input_data_set, θ), output_data)
 end
 
-function generate_data(ts, prob_set::Vector{ODEProblem}, isu0)
+function generate_data(ts, prob_set, isu0)
     batch_size = size(prob_set)[1]
     instances_size = size(ts)[2]
     dims = isu0 ? length(prob_set[1].u0) + 1 : length(prob_set[1].p) + 1
@@ -196,33 +227,10 @@ function generate_data(ts, prob_set::Vector{ODEProblem}, isu0)
     input_data_set
 end
 
-function generate_loss(
-        phi::PINOPhi{C, T, U}, train_set::TRAINSET, input_data_set, ts,
-        pino_phase::OperatorLearning) where {
-        C, T, U}
-    is_data_loss, is_physics_loss = pino_phase.is_data_loss, pino_phase.is_physics_loss
-    function loss(θ, _)
-        if is_data_loss
-            data_loss(phi, θ, train_set, input_data_set)
-        elseif is_physics_loss
-            physics_loss(phi, θ, ts, train_set, input_data_set)
-        elseif is_data_loss && is_physics_loss
-            data_loss(phi, θ, train_set, input_data_set) +
-            physics_loss(phi, θ, ts, train_set, input_data_set)
-        else
-            error("data loss or physics loss should be true")
-        end
-    end
-    return loss
-end
-
 function finetune_loss(phi::PINOPhi{C, T, U},
         θ,
-        train_set::TRAINSET,
         input_data_set,
         pino_phase::EquationSolving) where {C, T, U}
-    _, output_data = train_set.input_data, train_set.output_data
-    output_data = adapt(parameterless_type(ComponentArrays.getdata(θ)), output_data)
     pino_solution = pino_phase.pino_solution
     learned_operator = pino_solution.phi
     predict = learned_operator(input_data_set, pino_solution.res.u)
@@ -230,13 +238,44 @@ function finetune_loss(phi::PINOPhi{C, T, U},
 end
 
 function generate_loss(
-        phi::PINOPhi{C, T, U}, train_set::TRAINSET, input_data_set, ts,
-        pino_phase::EquationSolving) where {
+        phi::PINOPhi{C, T, U}, prob::ODEProblem, prob_set::Array, input_data_set, ts,
+        pino_phase::EquationSolving, isu0::Bool) where {
         C, T, U}
     a = 1 / 100
+    is_finetune_loss, is_physics_loss = pino_phase.is_finetune_loss,
+    pino_phase.is_physics_loss
     function loss(θ, _)
-        physics_loss(phi, θ, ts, train_set, input_data_set) +
-        a * finetune_loss(phi, θ, train_set, input_data_set, pino_phase)
+        if is_finetune_loss
+            finetune_loss(phi, θ, input_data_set, pino_phase)
+        elseif is_physics_loss
+            physics_loss(phi, prob, θ, ts, prob_set, input_data_set, isu0)
+        elseif is_finetune_loss && is_physics_loss
+            physics_loss(phi, prob, θ, ts, prob_set, input_data_set, isu0) +
+             a * finetune_loss(phi, θ, input_data_set, pino_phase)
+        else
+            error("finetune loss or physics loss should be true")
+        end
+    end
+    return loss
+end
+
+function generate_loss(
+        phi::PINOPhi{C, T, U},prob::ODEProblem, train_set::TRAINSET, input_data_set, ts,
+        pino_phase::OperatorLearning, isu0::Bool) where {
+        C, T, U}
+    is_data_loss, is_physics_loss = pino_phase.is_data_loss, pino_phase.is_physics_loss
+    prob_set, output_data = train_set.input_data, train_set.output_data
+    function loss(θ, _)
+        if is_data_loss
+            data_loss(phi, θ, output_data, input_data_set)
+        elseif is_physics_loss
+            physics_loss(phi, prob, θ, ts, prob_set, input_data_set, isu0)
+        elseif is_data_loss && is_physics_loss
+            data_loss(phi, θ, output_data, input_data_set) +
+            physics_loss(phi, prob, θ, ts, prob_set, input_data_set, isu0)
+        else
+            error("data loss or physics loss should be true")
+        end
     end
     return loss
 end
@@ -244,45 +283,40 @@ end
 function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem,
         alg::PINOODE,
         args...;
-        # dt = nothing,
         abstol = 1.0f-6,
         reltol = 1.0f-3,
         verbose = false,
         saveat = nothing,
         maxiters = nothing)
-    tspan = prob.tspan
+    @unpack tspan, u0, p, f = prob
     t0, t_end = tspan[1], tspan[2]
-    u0 = prob.u0
-    p = prob.p
-    # f = prob.f
-    # param_estim = alg.param_estim
-
-    chain = alg.chain
-    opt = alg.opt
-    init_params = alg.init_params
-    pino_phase = alg.pino_phase
-    # mapping between functional space of some vararible 'a' of equation (for example initial
-    # condition {u(t0 x)} or parameter p) and solution of equation u(t)
-    train_set = alg.train_set
+    @unpack chain, opt, pino_phase, init_params, isu0 = alg
 
     !(chain isa Lux.AbstractExplicitLayer) &&
         error("Only Lux.AbstractExplicitLayer neural networks are supported")
 
-    instances_size = size(train_set.output_data)[2]
-    range_ = range(t0, stop = t_end, length = instances_size)
-    ts = reshape(collect(range_), 1, instances_size)
-    prob_set, output_set = train_set.input_data, train_set.output_data
-    isu0 = train_set.isu0
-    input_data_set = generate_data(ts, prob_set, isu0)
-    # input_data_set =  if pino_phase == EquationSolving
-    #     generate_data(ts, [prob], isu0)
-    # elseif pino_phase == OperatorLearning
-    #     generate_data(ts, prob_set, isu0)
-    # else
-    #     error("pino_phase should be EquationSolving or OperatorLearning")
-    # end
-
-    if isu0 #TODO remove the block
+    if pino_phase isa EquationSolving
+        @unpack dt, pino_solution = pino_phase
+        range_ = collect(t0:dt:t_end)
+        ts = reshape(collect(range_), 1, size(range_)[1])
+        input_data_set = generate_data(ts, [prob], isu0)
+        if isu0
+            pino_solution.phi.u0 = input_data_set[2:end, :, :]
+        end
+    elseif pino_phase isa OperatorLearning
+        # mapping between functional space of some vararible 'a' of equation (for example initial
+        # condition {u(t0 x)} or parameter p) and solution of equation u(t)
+        train_set = pino_phase.train_set
+        instances_size = size(train_set.output_data)[2]
+        range_ = range(t0, stop = t_end, length = instances_size)
+        ts = reshape(collect(range_), 1, instances_size)
+        prob_set, _ = train_set.input_data, train_set.output_data
+        input_data_set = generate_data(ts, prob_set, isu0)
+    else
+        error("pino_phase should be EquationSolving or OperatorLearning")
+    end
+
+    if isu0
         u0 = input_data_set[2:end, :, :]
     else
         u0 = prob.u0
@@ -304,15 +338,10 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem,
     end
 
     if pino_phase isa EquationSolving
-        #TODO bad code rewrite,the parameter must uniquely match the index
-        #TODO doenst need TRAINSET for EquationSolving
-        find(as, a) = findfirst(x -> isapprox(x.p, a.p), as)
-        index = find(prob_set, prob)
-        input_data_set = input_data_set[:, :, [index]]
-        train_set = TRAINSET(prob_set[index:index], output_set[:, :, [index]], isu0)
-        total_loss = generate_loss(phi, train_set, input_data_set, ts, pino_phase)
+        prob_set = [prob]
+        total_loss = generate_loss(phi, prob, prob_set, input_data_set, ts, pino_phase,isu0)
     elseif pino_phase isa OperatorLearning
-        total_loss = generate_loss(phi, train_set, input_data_set, ts, pino_phase)
+        total_loss = generate_loss(phi, prob, train_set, input_data_set, ts, pino_phase, isu0)
     else
         error("pino_phase should be EquationSolving or OperatorLearning")
     end
diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl
index 9c53e528a..628e68c99 100644
--- a/test/PINO_ode_tests.jl
+++ b/test/PINO_ode_tests.jl
@@ -10,7 +10,9 @@ using NeuralPDE
     linear = (u, p, t) -> cos(p * t)
     tspan = (0.0f0, 2.0f0)
     u0 = 0.0f0
-    #generate data set
+    p = pi / 2
+    prob = ODEProblem(linear, u0, tspan, p)
+    #generate data
     t0, t_end = tspan
     instances_size = 50
     range_ = range(t0, stop = t_end, length = instances_size)
@@ -21,58 +23,65 @@ using NeuralPDE
     u_output_ = zeros(Float32, 1, instances_size, batch_size)
     prob_set = []
     for (i, a_i) in enumerate(as)
-        prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i)
-        sol1 = solve(prob, Tsit5(); saveat = 0.0204)
+        prob_ = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i)
+        sol1 = solve(prob_, Tsit5(); saveat = 0.0204)
         reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1))
-        push!(prob_set, prob)
+        push!(prob_set, prob_)
         u_output_[:, :, i] = reshape_sol
     end
-
     """
-    Set of training data:
-    * input data: set of parameters 'a'
+    Training data:
+    * input data: set of parameters 'a',
     * output data: set of solutions u(t){a} corresponding parameter 'a'.
-     """
+    """
     train_set = TRAINSET(prob_set, u_output_)
-    p = pi / 2
-    prob = ODEProblem(linear, u0, tspan, p)
+
+    # operator learning phase
+    # init neural network
     chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ),
-        Lux.Dense(16, 16, Lux.σ),
         Lux.Dense(16, 16, Lux.σ),
         Lux.Dense(16, 16, Lux.σ),
         Lux.Dense(16, 32, Lux.σ),
+        Lux.Dense(32, 32, Lux.σ),
         Lux.Dense(32, 1))
     # flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,),
     #     σ = gelu)
-    opt = OptimizationOptimisers.Adam(0.03)
-    # pino_phase = OperatorLearning(train_set, is_data_loss = true, is_physics_loss = true)
-    pino_phase = OperatorLearning(
-        is_data_loss = true, is_physics_loss = true)
-    alg = PINOODE(chain, opt, train_set, pino_phase)
-    # pino_solution = learn()
-    pino_solution = solve(prob, alg, verbose = false, maxiters = 2000)
+
+    opt = OptimizationOptimisers.Adam(0.01)
+    pino_phase = OperatorLearning(train_set, is_data_loss = true, is_physics_loss = true)
+
+    alg = PINOODE(chain, opt, pino_phase)
+    pino_solution = solve(
+        prob, alg, verbose = false, maxiters = 2000)
     predict = pino_solution.predict
     ground = u_output_
     @test ground≈predict atol=1
+    # equation solving phase
+    dt = (t_end - t0) / instances_size
+    pino_phase = EquationSolving(dt, pino_solution)
+    chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ),
+        Lux.Dense(16, 16, Lux.σ),
+        Lux.Dense(16, 32, Lux.σ),
+        Lux.Dense(32, 1))
+    alg = PINOODE(chain, opt, pino_phase)
+    fine_tune_solution = solve(
+        prob, alg, verbose = false, maxiters = 2000)
 
-    pino_phase = EquationSolving(pino_solution)
-    alg = PINOODE(chain, opt, train_set, pino_phase)
-    pino_solution = solve(prob, alg, verbose = true, maxiters = 2000)
-
-    find(as, a) = findfirst(x -> isapprox(x.p, a.p), as)
-    index = find(prob_set, prob)
-    predict = pino_solution.predict
-    ground = u_output_[:,:, [index]]
-    @test ground≈predict atol=0.1
+    fine_tune_predict = fine_tune_solution.predict
+    operator_predict = pino_solution.phi(
+        fine_tune_solution.input_data_set, pino_solution.res.u)
+    ground = linear_analytic.(u0, p, fine_tune_solution.input_data_set[[1], :, :])
+    @test ground≈fine_tune_predict atol=0.1
+    @test operator_predict≈fine_tune_predict rtol=0.1
 end
 
-
 @testset "Example u0" begin
     linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p)
     linear = (u, p, t) -> cos(p * t)
     tspan = (0.0f0, 2.0f0)
     p = Float32(pi)
     u0 = 2.0f0
+    prob = ODEProblem(linear, u0, tspan, p)
     #generate data set
     t0, t_end = tspan
     instances_size = 50
@@ -83,33 +92,48 @@ end
     u_output_ = zeros(Float32, 1, instances_size, batch_size)
     prob_set = []
     for (i, u0_i) in enumerate(u0s)
-        prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0_i, tspan, p)
-        sol1 = solve(prob, Tsit5(); saveat = 0.0204)
+        prob_ = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0_i, tspan, p)
+        sol1 = solve(prob_, Tsit5(); saveat = 0.0204)
         reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1))
-        push!(prob_set, prob)
+        push!(prob_set, prob_)
         u_output_[:, :, i] = reshape_sol
     end
-
+    # operator learning phase
     """
       Set of training data:
       * input data: set of initial conditions 'u0'
       * output data: set of solutions u(t){u0} corresponding initial conditions 'u0'.
     """
-    train_set = TRAINSET(prob_set, u_output_; isu0 = true)
-    #TODO we argument u0 but dont actually use u0 because we use only set of u0 for generate train set from prob_set
-    prob = ODEProblem(linear, 0.0f0, tspan, p)
+    train_set = TRAINSET(prob_set, u_output_)
     # fno = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu)
     chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ),
-        Lux.Dense(16, 16, Lux.σ),
         Lux.Dense(16, 16, Lux.σ),
         Lux.Dense(16, 16, Lux.σ),
         Lux.Dense(16, 32, Lux.σ),
+        Lux.Dense(32, 32, Lux.σ),
         Lux.Dense(32, 1))
     opt = OptimizationOptimisers.Adam(0.001)
-    pino_phase = OperatorLearning()
-    alg = PINOODE(chain, opt, train_set, pino_phase)
+    pino_phase = OperatorLearning(train_set)
+    alg = PINOODE(chain, opt, pino_phase; isu0 = true)
     pino_solution = solve(prob, alg, verbose = false, maxiters = 2000)
     predict = pino_solution.predict
     ground = u_output_
-    @test ground≈predict atol=1.0
+    @test ground≈predict atol=1.
+
+    # equation solving phase
+    dt = (t_end -t0) / instances_size
+    pino_phase = EquationSolving(dt, pino_solution)
+    chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ),
+        Lux.Dense(16, 16, Lux.σ),
+        Lux.Dense(16, 32, Lux.σ),
+        Lux.Dense(32, 1))
+    alg = PINOODE(chain, opt, pino_phase; isu0 = true)
+    fine_tune_solution = solve(prob, alg, verbose = false, maxiters = 2000)
+
+    fine_tune_predict = fine_tune_solution.predict
+    operator_predict = pino_solution.phi(
+        fine_tune_solution.input_data_set, pino_solution.res.u)
+    ground_fine_tune = linear_analytic.(u0, p, fine_tune_solution.input_data_set[[1], :, :])
+    @test ground_fine_tune≈fine_tune_predict atol=1
+    @test operator_predict≈fine_tune_predict rtol=0.1
 end
diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl
index 378ce9ddd..919e6012e 100644
--- a/test/PINO_ode_tests_gpu.jl
+++ b/test/PINO_ode_tests_gpu.jl
@@ -16,21 +16,23 @@ const gpud = gpu_device()
     linear = (u, p, t) -> cos(p * t)
     tspan = (0.0f0, 2.0f0)
     u0 = 0.0f0
+    p = pi / 2.0f0
+    prob = ODEProblem(linear, u0, tspan, p)
     #generate data set
     t0, t_end = tspan
     instances_size = 50
     range_ = range(t0, stop = t_end, length = instances_size)
     ts = reshape(collect(range_), 1, instances_size)
     batch_size = 50
-    as = [Float32(i) for i in range(0.1, stop = pi / 2, length = batch_size)]
+    as = [Float32(i) for i in range(0.1, stop = pi / 2.0f0, length = batch_size)]
 
     u_output_ = zeros(Float32, 1, instances_size, batch_size)
     prob_set = []
     for (i, a_i) in enumerate(as)
-        prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i)
-        sol1 = solve(prob, Tsit5(); saveat = 0.0204)
+        prob_ = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i)
+        sol1 = solve(prob_, Tsit5(); saveat = 0.0204)
         reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1))
-        push!(prob_set, prob)
+        push!(prob_set, prob_)
         u_output_[:, :, i] = reshape_sol
     end
     u_output_ = u_output_ |> gpud
@@ -42,7 +44,6 @@ const gpud = gpu_device()
     """
     train_set = TRAINSET(prob_set, u_output_)
 
-    prob = ODEProblem(linear, u0, tspan, 0)
     inner = 50
     chain = Lux.Chain(Lux.Dense(2, inner, Lux.σ),
         Lux.Dense(inner, inner, Lux.σ),
@@ -52,12 +53,26 @@ const gpud = gpu_device()
         Lux.Dense(inner, 1))
     ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud
     opt = OptimizationOptimisers.Adam(0.03)
-    pino_phase = OperatorLearning()
-    alg = PINOODE(chain, opt, train_set, pino_phase; init_params = ps)
+    pino_phase = OperatorLearning(train_set)
+    alg = PINOODE(chain, opt, pino_phase; init_params = ps)
     pino_solution = solve(prob, alg, verbose = false, maxiters = 2000)
     predict = pino_solution.predict |> cpu
     ground = u_output_ |> cpu
     @test ground≈predict atol=1
+
+    dt = (t_end - t0) / instances_size
+    pino_phase = EquationSolving(dt, pino_solution)
+    alg = PINOODE(chain, opt, pino_phase; init_params = ps)
+    fine_tune_solution = solve(
+        prob, alg, verbose = false, maxiters = 2000)
+
+    fine_tune_predict = fine_tune_solution.predict |> cpu
+    operator_predict = pino_solution.phi(
+        fine_tune_solution.input_data_set, pino_solution.res.u) |> cpu
+    input_data_set_ = fine_tune_solution.input_data_set[[1], :, :] |> cpu
+    ground_fine_tune = linear_analytic.(u0, p, input_data_set_)
+    @test ground_fine_tune≈fine_tune_predict atol=0.5
+    @test operator_predict≈fine_tune_predict rtol=0.1
 end
 
 @testset "lotka volterra" begin
@@ -76,26 +91,26 @@ end
     p = Float32[1.5, 1.0, 3.0, 1.0]
     tspan = (0.0f0, 4.0f0)
     dt = 0.01f0
+    prob = ODEProblem(lotka_volterra, u0, tspan, p)
     t0, t_end = tspan
-
     instances_size = 100
     range_ = range(t0, stop = t_end, length = instances_size)
     ts = reshape(collect(range_), 1, instances_size)
     batch_size = 50
-    ps = [p .+ i * Float32[0.000, 0.0, 0.001, 0.01] for i in 1:batch_size]
+    ps = [p .+ (i - 1) * Float32[0.000, 0.0, 0.001, 0.01] for i in 1:batch_size]
     u_output_ = zeros(Float32, 2, instances_size, batch_size)
     prob_set = []
     for (i, p_i) in enumerate(ps)
-        prob = ODEProblem(lotka_volterra, u0, tspan, p_i)
-        solution = solve(prob, Tsit5(); saveat = dt)
+        prob_ = ODEProblem(lotka_volterra, u0, tspan, p_i)
+        solution = solve(prob_, Tsit5(); saveat = dt)
         reshape_sol_ = reduce(hcat, solution(range_).u)
         reshape_sol = Float32.(reshape(reshape_sol_, 2, instances_size, 1))
-        push!(prob_set, prob)
+        push!(prob_set, prob_)
         u_output_[:, :, i] = reshape_sol
     end
 
     train_set = TRAINSET(prob_set, u_output_)
-    prob = ODEProblem(lotka_volterra, u0, tspan, p)
+
     # flat_no = FourierNeuralOperator(ch = (5, 64, 64, 64, 64, 64, 128, 2), modes = (16,),
     #     σ = gelu)
     # flat_no = Lux.transform(flat_no)
@@ -109,11 +124,28 @@ end
     ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud
 
     opt = OptimizationOptimisers.Adam(0.001)
-    pino_phase = OperatorLearning()
-    alg = PINOODE(
-        chain, opt, train_set, pino_phase; init_params = ps)
+    pino_phase = OperatorLearning(train_set)
+    alg = PINOODE(chain, opt, pino_phase; init_params = ps)
     pino_solution = solve(prob, alg, verbose = false, maxiters = 4000)
     predict = pino_solution.predict |> cpu
     ground = u_output_
     @test ground≈predict atol=5
+
+    dt = (t_end - t0) / instances_size
+    pino_phase = EquationSolving(dt, pino_solution; is_finetune_loss = true,is_physics_loss = true)
+    chain = Lux.Chain(Lux.Dense(5, 16, Lux.σ),
+        Lux.Dense(16, 16, Lux.σ),
+        Lux.Dense(16, 32, Lux.σ),
+        Lux.Dense(32, 2))
+    ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud
+    alg = PINOODE(chain, opt, pino_phase; init_params = ps)
+    fine_tune_solution = solve(prob, alg, verbose = false, maxiters = 2000)
+
+    fine_tune_predict = fine_tune_solution.predict |> cpu
+    operator_predict = pino_solution.phi(
+        fine_tune_solution.input_data_set, pino_solution.res.u) |> cpu
+    input_data_set_ = fine_tune_solution.input_data_set[[1], :, :] |> cpu
+    ground_fine_tune = u_output_[:, :, [1]]
+    @test ground_fine_tune ≈ fine_tune_predict[:, 1:100, :] atol = 3
+    @test operator_predict≈fine_tune_predict rtol=0.1
 end