diff --git a/docs/pages.jl b/docs/pages.jl
index 020d330407..54b706f28c 100644
--- a/docs/pages.jl
+++ b/docs/pages.jl
@@ -32,6 +32,7 @@ pages = ["index.md",
         "manual/training_strategies.md",
         "manual/adaptive_losses.md",
         "manual/logging.md",
-        "manual/neural_adapters.md"],
+        "manual/neural_adapters.md",
+        "manual/pino_ode.md"],
     "Developer Documentation" => Any["developer/debugging.md"]
 ]
diff --git a/docs/src/manual/pino_ode.md b/docs/src/manual/pino_ode.md
index 379f84c46b..72b5385924 100644
--- a/docs/src/manual/pino_ode.md
+++ b/docs/src/manual/pino_ode.md
@@ -5,17 +5,7 @@ PINOODE
 ```
 
 ```@docs
-TRAINSET
+DeepONet
 ```
 
-```@docs
-PINOsolution
-```
 
-```@docs
-OperatorLearning
-```
-
-```@docs
-EquationSolving
-```
diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md
index 2e0ab0ed43..2d6d053852 100644
--- a/docs/src/tutorials/pino_ode.md
+++ b/docs/src/tutorials/pino_ode.md
@@ -1,8 +1,6 @@
-#  Physics informed Neural Operator ODEs Solvers
+# Physics Informed Neural Operator for ODEs Solvers
 
-This tutorial is an introduction to using physics-informed neural operator (PINOs) for solving family of parametric ordinary diferential equations (ODEs). 
-
-#TODO two phase
+This tutorial provides an example of how using Physics Informed Neural Operator (PINO) for solving family of parametric ordinary differential equations (ODEs).
 
 ## Operator Learning  for a family of parametric ODE.
 
@@ -11,95 +9,62 @@ using Test
 using OrdinaryDiffEq, OptimizationOptimisers
 using Lux
 using Statistics, Random
-# 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)
+equation = (u, p, t) -> cos(p * t)
+tspan = (0.0f0, 1.0f0)
+u0 = 1.0f0
+prob = ODEProblem(equation,  u0, tspan)
+
+# initilize DeepONet operator
+branch = Lux.Chain(
+    Lux.Dense(1, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10))
+trunk = Lux.Chain(
+    Lux.Dense(1, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10, Lux.tanh_fast))
+
+deeponet = NeuralPDE.DeepONet(branch, trunk; linear = nothing)
+
+bounds = (p = [0.1f0, pi],)
+#TODO add truct
+strategy  = (branch_size = 50, trunk_size = 40)
+# strategy = (branch_size = 50, dt = 0.1)?
+opt = OptimizationOptimisers.Adam(0.03)
+alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy)
+
+sol = solve(prob, alg, verbose = true, maxiters = 2000)
 ```
 
-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:
-* input data: set of parameters 'a',
-* output data: set of solutions u(t){a} corresponding parameter 'a'.
+Now let's compare the prediction from the learned operator with the ground truth solution which is obtained by analytic solution the parametric ODE. Where 
+Compare prediction with ground truth.
 
 ```@example pino
-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)]
-
-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)
-    reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1))
-    push!(prob_set, prob_)
-    u_output_[:, :, i] = reshape_sol
-end
-train_set = TRAINSET(prob_set, u_output_)
+using Plots
+# Compute the ground truth solution for each parameter value and time in the solution
+# The '.' operator is used to apply the functd ion element-wise
+ground_analytic = (u0, p, t) -> begin u0 + sin(p * t) / (p)
+p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size)
+p = reshape(p_, 1, branch_size, 1)
+ground_solution = ground_analytic.(u0, p, sol.t.trunk)
+
+# Plot the predicted solution and the ground truth solution as a filled contour plot
+# sol.u[1, :, :], represents the predicted solution for each parameter value and time
+plot(sol.u[1, :, :], linetype = :contourf)
+plot!(ground_solution[1, :, :], linetype = :contourf)
 ```
 
-Here it used the PINO method to learning operator of the given family of parametric ODEs. 
 
 ```@example pino
-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_
-```
+using Plots
 
-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.
+# 'i' is the index of the parameter 'a' in the dataset
+i = 45
 
-```@example pino
-using Plots
-i=45
+# 'predict' is the predicted solution from the PINO model
+# 'ground' is the ground truth solution
 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/NeuralPDE.jl b/src/NeuralPDE.jl
index d9f401151d..30ac9b32d3 100644
--- a/src/NeuralPDE.jl
+++ b/src/NeuralPDE.jl
@@ -59,7 +59,7 @@ include("PDE_BPINN.jl")
 include("dgm.jl")
 
 
-export NNODE, NNDAE, PINOODE, DeepONet
+export NNODE, NNDAE, PINOODE, DeepONet, SomeStrategy #TODO remove SomeStrategy
        PhysicsInformedNN, discretize,
        GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining,
        WeightedIntervalTraining,
diff --git a/src/neural_operators.jl b/src/neural_operators.jl
index 7544675832..ddc87dcfce 100644
--- a/src/neural_operators.jl
+++ b/src/neural_operators.jl
@@ -1,47 +1,102 @@
-#TODO: Add docstrings
+abstract type NeuralOperator <: Lux.AbstractExplicitLayer end
+
 """
 DeepONet(branch,trunk)
 """
-struct DeepONet{} <: Lux.AbstractExplicitLayer
+
+"""
+    DeepONet(branch,trunk,linear=nothing)
+
+`DeepONet` is differential neural operator focused for solving physic-informed parametric ODEs.
+
+DeepONet uses two neural networks, referred to as the "branch" and "trunk", to approximate
+the solution of a differential equation. The branch network takes the spatial variables as
+input and the trunk network takes the temporal variables as input. The final output is
+the dot product of the outputs of the branch and trunk networks.
+
+DeepONet is composed of two separate neural networks referred to as the "branch" and "trunk",
+respectively. The branch net takes on input represents a function  evaluated at a collection
+of fixed locations in some boundsand returns a features embedding. The trunk net takes the
+continuous coordinates as inputs, and outputs a features embedding. The final output of the
+DeepONet, the outputs of the branch and trunk networks are merged together via a dot product.
+
+## Positional Arguments
+*  `branch`: A branch neural network.
+*  `trunk`: A trunk neural network.
+
+## Keyword Arguments
+* `linear`: A linear layer to apply to the output of the branch and trunk networks.
+
+## Example
+
+```julia
+branch = Lux.Chain(
+    Lux.Dense(1, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10))
+trunk = Lux.Chain(
+    Lux.Dense(1, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10, Lux.tanh_fast))
+linear = Lux.Chain(Lux.Dense(10, 1))
+
+deeponet = DeepONet(branch, trunk; linear= linear)
+
+a = rand(1, 50, 40)
+b = rand(1, 1, 40)
+x = (branch = a, trunk = b)
+θ, st = Lux.setup(Random.default_rng(), deeponet)
+y, st = deeponet(x, θ, st)
+```
+
+## References
+* Lu Lu, Pengzhan Jin, George Em Karniadakis "DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators"
+* Sifan Wang "Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets"
+"""
+
+struct DeepONet{L <: Union{Nothing, Lux.AbstractExplicitLayer }} <: NeuralOperator
     branch::Lux.AbstractExplicitLayer
     trunk::Lux.AbstractExplicitLayer
+    linear::L
+end
+
+function DeepONet(branch, trunk; linear=nothing)
+    DeepONet(branch, trunk, linear)
 end
 
 function Lux.setup(rng::AbstractRNG, l::DeepONet)
-    branch, trunk = l.branch, l.trunk
+    branch, trunk, linear = l.branch, l.trunk, l.linear
     θ_branch, st_branch = Lux.setup(rng, branch)
     θ_trunk, st_trunk = Lux.setup(rng, trunk)
     θ = (branch = θ_branch, trunk = θ_trunk)
     st = (branch = st_branch, trunk = st_trunk)
+    if linear !== nothing
+        θ_liner, st_liner = Lux.setup(rng, linear)
+        θ = (θ..., liner = θ_liner)
+        st = (st..., liner = st_liner)
+    end
     θ, st
 end
 
-# function Lux.initialparameters(rng::AbstractRNG, e::DeepONet)
-#     code
-# end
-
 Lux.initialstates(::AbstractRNG, ::DeepONet) = NamedTuple()
 
-"""
-example:
-
-branch = Lux.Chain(Lux.Dense(1, 32, Lux.σ), Lux.Dense(32, 1))
-trunk = Lux.Chain(Lux.Dense(1, 32, Lux.σ), Lux.Dense(32, 1))
-a = rand(1, 100, 10)
-t = rand(1, 1, 10)
-x = (branch = a, trunk = t)
-
-deeponet = DeepONet(branch, trunk)
-θ, st = Lux.setup(Random.default_rng(), deeponet)
-y = deeponet(x, θ, st)
-"""
 @inline function (f::DeepONet)(x::NamedTuple, θ, st::NamedTuple)
-    parameters, cord = x.branch, x.trunk
+    x_branch, x_trunk = x.branch, x.trunk
     branch, trunk = f.branch, f.trunk
     st_branch, st_trunk = st.branch, st.trunk
     θ_branch, θ_trunk = θ.branch, θ.trunk
-    out_b, st_b = branch(parameters, θ_branch, st_branch)
-    out_t, st_t = trunk(cord, θ_trunk, st_trunk)
-    out = out_b' * out_t
-    return out, (branch = st_b, trunk = st_t)
+    out_b, st_b = branch(x_branch, θ_branch, st_branch)
+    out_t, st_t = trunk(x_trunk, θ_trunk, st_trunk)
+    if f.linear !== nothing
+        linear = f.linear
+        θ_liner, st_liner = θ.liner, st.liner
+        # out = sum(out_b .* out_t, dims = 1)
+        out_ = out_b .* out_t
+        out, st_liner = linear(out_, θ_liner, st_liner)
+        out = sum(out, dims = 1)
+        return out, (branch = st_b, trunk = st_t, liner = st_liner)
+    else
+        out = sum(out_b .* out_t, dims = 1)
+        return out, (branch = st_b, trunk = st_t)
+    end
 end
diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl
index 597131ee0d..09d4fdefc4 100644
--- a/src/pino_ode_solve.jl
+++ b/src/pino_ode_solve.jl
@@ -1,36 +1,42 @@
 """
-   PINOODE(chain,
+ PINOODE
+
+PINOODE(chain,
     OptimizationOptimisers.Adam(0.1),
-    pino_phase;
-    init_params,
+    bounds;
+    init_params = nothing,
+    strategy = nothing
     kwargs...)
 
-The method is that combine training data and physics constraints
-to learn the solution operator of a given family of parametric Ordinary Differential Equations (ODE).
+Algorithm for solving paramentric ordinary differential equations using a physics-informed
+neural operator, which is used as a solver for a parametrized `ODEProblem`.
 
 ## Positional Arguments
 * `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.
-* `pino_phase`: The phase of the PINN algorithm, either `OperatorLearning` or `EquationSolving`.
+* `bounds`: A dictionary containing the bounds for the parameters of the neural network
+in which will be train the prediction of parametric ODE.
 
 ## Keyword Arguments
 * `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'.
+* `strategy`: The strategy for training the neural network.
+* `additional_loss`: additional function to the main one. For example, add training on data.
 * `kwargs`: Extra keyword arguments are splatted to the Optimization.jl `solve` call.
 
 ## References
 * Sifan Wang "Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets"
 * Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations"
 """
-struct PINOODE{C, O, B, I, S, K} <: SciMLBase.AbstractODEAlgorithm
+struct PINOODE{C, O, B, I, S <: Union{Nothing, AbstractTrainingStrategy}, AL <: Union{Nothing, Function},K} <:
+       SciMLBase.AbstractODEAlgorithm
     chain::C
     opt::O
     bounds::B
     init_params::I
-    isu0::Bool
     strategy::S
+    additional_loss::AL #TODO
     kwargs::K
 end
 
@@ -38,17 +44,26 @@ function PINOODE(chain,
         opt,
         bounds;
         init_params = nothing,
-        isu0 = false, #TODOD remove
         strategy = nothing,
+        additional_loss = nothing,
         kwargs...)
     !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain))
-    PINOODE(chain, opt, bounds, init_params, isu0, strategy, kwargs)
+    PINOODE(chain, opt, bounds, init_params, strategy, additional_loss, kwargs)
+end
+
+#TODO cool name for strategy
+struct SomeStrategy{} <: AbstractTrainingStrategy
+    branch_size::Int
+    trunk_size::Int
+end
+function SomeStrategy(; branch_size = 10, trunk_size=10)
+    SomeStrategy(branch_size, trunk_size)
 end
 
 mutable struct PINOPhi{C, T, U, S}
     chain::C
     t0::T
-    u0::U
+    u0::U#TODO remove u0, t0?
     st::S
     function PINOPhi(chain::Lux.AbstractExplicitLayer, t0, u0, st)
         new{typeof(chain), typeof(t0), typeof(u0), typeof(st)}(chain, t0, u0, st)
@@ -68,100 +83,99 @@ function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer,
     PINOPhi(chain, t0, u0, st), init_params
 end
 
-#TODO update
-# function (f::PINOPhi{C, T, U})(t::AbstractArray,
-#         θ) where {C <: Lux.AbstractExplicitLayer, T, U}
-#     y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st)
-#     ChainRulesCore.@ignore_derivatives f.st = st
-#     ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), t[1:size(y)[1], :, :])
-#     u_0 = adapt(parameterless_type(ComponentArrays.getdata(θ)), f.u0)
-#     u_0 .+ (ts .- f.t0) .* y
-# end
-
-#TODO C <: DeepONet
 function (f::PINOPhi{C, T, U})(x::NamedTuple, θ) where {C, T, U}
     y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st)
     ChainRulesCore.@ignore_derivatives f.st = st
-    a, t = x.branch, x.trunk
-    ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), t)
-    u0_ = adapt(parameterless_type(ComponentArrays.getdata(θ)), f.u0)
-    u0_ .+ (ts .- f.t0) .* y
+    y
 end
 
-# function dfdx(phi::PINOPhi, t::AbstractArray, θ)
-#     ε = [sqrt(eps(eltype(t))), zeros(eltype(t), size(t)[1] - 1)...]
-#     (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t)))
-# end
-
-#TODO C <: DeepONet
-function dfdx(phi::PINOPhi{C, T, U}, x::NamedTuple, θ) where {C, T, U}
-    t = x.trunk
-    ε = [sqrt(eps(eltype(t)))]
-    phi_trunk(x, θ) = phi.chain.trunk(x, θ.trunk, phi.st.trunk)[1]
-    du_trunk_ = (phi_trunk(t .+ ε, θ) .- phi_trunk(t, θ)) ./ sqrt(eps(eltype(t)))
-    u_branch = phi.chain.branch(x.branch, θ.branch, phi.st.branch)[1]
-    u_branch' .* du_trunk_
+function dfdx(phi::PINOPhi{C, T, U}, x::NamedTuple, θ, branch_left) where {C, T, U}
+    x_trunk = x.trunk
+    x_left = (branch = branch_left, trunk = x_trunk .+ sqrt(eps(eltype(x_trunk))))
+    x_right = (branch = x.branch, trunk = x_trunk)
+    (phi(x_left, θ) .- phi(x_right, θ)) / sqrt(eps(eltype(x_trunk)))
 end
 
-# function l₂loss(𝐲̂, 𝐲)
-#     feature_dims = 2:(ndims(𝐲) - 1)
-#     loss = sum(.√(sum(abs2, 𝐲̂ - 𝐲, dims = feature_dims)))
-#     y_norm = sum(.√(sum(abs2, 𝐲, dims = feature_dims)))
-#     return loss / y_norm
-# end
-
 function physics_loss(phi::PINOPhi{C, T, U}, prob::ODEProblem, x, θ) where {C, T, U}
+    x_ , p = x
+    norm = prod(size(x_.branch))
+    branch_left = prob.f.(nothing, p, x_.trunk .+ sqrt(eps(eltype(x_.trunk))))
+    sum(abs2, vec(dfdx(phi, x_, θ, branch_left)) .- vec(x_.branch)) / norm
+end
+
+function operator_loss(phi::PINOPhi{C, T, U}, prob::ODEProblem, x, θ) where {C, T, U}
+    x_, u0_ = x
+    norm = prod(size(u0_))
+    sum(abs2, vec(phi(x_, θ)) .- vec(u0_)) / norm
+end
+
+#TODO inside loss couse call f(),return (p, t)
+function get_trainset(branch_size, trunk_size, bounds, tspan, prob)
+    p_ = range(bounds.p[1], stop = bounds.p[2], length = branch_size)
+    p = reshape(p_, 1, branch_size,1)
+    t_ = collect(range(tspan[1], stop = tspan[2], length = trunk_size))
+    t = reshape(t_, 1, 1, trunk_size)
     f = prob.f
-    ps , ts  = x.branch, x.trunk
-    norm = size(x.branch)[2] * size(x.trunk)[2]
-    sum(abs2, dfdx(phi, x, θ) - f.(phi(x, θ), ps, ts)) / norm
+    #TODO phi
+    f_ = f.(nothing, p, t)
+    x = (branch = f_, trunk = t)
+    x , p
 end
 
-function get_trainset(bounds, tspan, strategy)
-    #TODO dt -> instances_size
-    instances_size = 100
-    p = range(bounds.p[1], stop = bounds.p[2], length = instances_size)
-    t = range(tspan[1], stop = tspan[2], length = instances_size)
-    x = (branch = collect(p)', trunk = collect(t)')
-    x
+#TODO return (t0, u0)?
+function get_trainset_operator(branch_size, trunk_size, bounds, tspan, prob)
+    p_ = range(bounds.p[1], stop = bounds.p[2], length = branch_size)
+    p = reshape(p_, 1, branch_size, 1)
+    t_ = collect(range(tspan[1], stop = tspan[2], length = trunk_size))
+    t = reshape(t_, 1, 1, trunk_size)
+    f = prob.f
+    #TODO phi
+    f_ = f.(nothing, p, t[1])
+    x_ = (branch = f_, trunk = t)
+
+    t0 = x_.trunk[:, :, [1]]
+    # f_2 = f.(nothing, p, t0)
+    f_2 = f_[:, :, [1]]
+    x = (branch = f_2, trunk =t0)
+    u0_ = fill(prob.u0, size(f_2))
+    x, u0_
 end
 
-#TODO GridTraining
 function generate_loss(strategy, prob::ODEProblem, phi, bounds, tspan)
-    x = get_trainset(bounds, tspan, strategy)
+    branch_size, trunk_size = strategy.branch_size, strategy.trunk_size
+    #TODO move to loss
+    x = get_trainset(branch_size, trunk_size, bounds, tspan, prob)
+    x_op = get_trainset_operator(branch_size, trunk_size, bounds, tspan, prob)
     function loss(θ, _)
-        physics_loss(phi, prob, x, θ)
+        physics_loss(phi, prob, x, θ) + operator_loss(phi, prob, x_op, θ)
     end
 end
 
 function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem,
         alg::PINOODE,
         args...;
-        dt = nothing,
-        abstol = 1.0f-6,
+        abstol = 1.0f-8,
         reltol = 1.0f-3,
         verbose = false,
         saveat = nothing,
         maxiters = nothing)
     @unpack tspan, u0, p, f = prob
-    t0, t_end = tspan[1], tspan[2]
-    @unpack chain, opt, bounds, init_params, isu0 = alg
+    @unpack chain, opt, bounds, init_params, strategy = alg
 
     !(chain isa Lux.AbstractExplicitLayer) &&
         error("Only Lux.AbstractExplicitLayer neural networks are supported")
 
     if !any(in(keys(bounds)), (:u0, :p))
-        error("bounds should contain u0 or p only")
+        error("bounds should contain u0 and p only")
     end
 
-    phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params)
-    init_params = ComponentArrays.ComponentArray(init_params)
+    phi, init_params = generate_pino_phi_θ(chain, 0, u0, init_params)
 
     isinplace(prob) &&
         throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t)."))
 
     try
-        x = (branch = rand(length(bounds), 10), trunk = rand(1, 10))
+        x = (branch = rand(1, 10,10), trunk = rand(1, 1, 10))
         phi(x, init_params)
     catch err
         if isa(err, DimensionMismatch)
@@ -171,13 +185,12 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem,
         end
     end
 
-    strategy = nothing
-
     inner_f = generate_loss(strategy, prob, phi, bounds, tspan)
 
+    #TODO impl add loss
     function total_loss(θ, _)
         inner_f(θ, nothing)
-        #TODO add loss
+
         # L2_loss = inner_f(θ, nothing)
         # if !(additional_loss isa Nothing)
         #     L2_loss = L2_loss + additional_loss(phi, θ)
@@ -199,16 +212,19 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem,
 
     optprob = OptimizationProblem(optf, init_params)
     res = solve(optprob, opt; callback, maxiters, alg.kwargs...)
-    res, phi
-
-    #TODO build_solution
-    # if saveat isa Number
-    #     ts = tspan[1]:saveat:tspan[2]
-    # end
-
-    # if u0 isa Number
-    #     u = [first(phi(t, res.u)) for t in ts]
-    # end
-    # sol = SciMLBase.build_solution(prob, alg, ts, u;
-    # sol
+
+    branch_size, trunk_size = strategy.branch_size, strategy.trunk_size
+    x, p = get_trainset(branch_size, trunk_size, bounds, tspan, prob)
+    u = phi(x, res.u)
+
+    sol = SciMLBase.build_solution(prob, alg, x, u;
+        k = res, dense = true,
+        calculate_error = false,
+        retcode = ReturnCode.Success,
+        original = res,
+        resid = res.objective)
+    SciMLBase.has_analytic(prob.f) &&
+        SciMLBase.calculate_solution_errors!(sol; timeseries_errors = true,
+            dense_errors = false)
+    sol
 end
diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl
index 50efbe5b43..5049c121d7 100644
--- a/test/PINO_ode_tests.jl
+++ b/test/PINO_ode_tests.jl
@@ -4,101 +4,83 @@ using Lux
 using Statistics, Random
 using NeuralPDE
 
-@testset "Example p" begin
+# dG(u(t,p),t) = u(t,p)
+@testset "Example du = cos(p * t)" begin
     equation = (u, p, t) -> cos(p * t)
-    tspan = (0.0f0, 2.0f0)
-    u0 = 0.0f0
-    prob = ODEProblem(equation, u0, tspan)
-    # prob = PINOODEProblem(equation, tspan)?
+    tspan = (0.0f0, 1.0f0)
+    u0 = 1.0f0
+    prob = ODEProblem(equation,  u0, tspan)
+
+    branch = Lux.Chain(
+        Lux.Dense(1, 10, Lux.tanh_fast),
+        Lux.Dense(10, 10, Lux.tanh_fast),
+        Lux.Dense(10, 10))
+    trunk = Lux.Chain(
+        Lux.Dense(1, 10, Lux.tanh_fast),
+        Lux.Dense(10, 10, Lux.tanh_fast),
+        Lux.Dense(10, 10, Lux.tanh_fast))
+
+    deeponet = NeuralPDE.DeepONet(branch, trunk; linear = nothing)
+    a = rand(1, 50, 40)
+    b = rand(1, 1, 40)
+    x = (branch = a, trunk = b)
+    θ, st = Lux.setup(Random.default_rng(), deeponet)
+    c = deeponet(x, θ, st)[1]
 
-    # init neural operator
-    branch = Lux.Chain(Lux.Dense(1, 32, Lux.σ), Lux.Dense(32, 32, Lux.σ), Lux.Dense(32, 1))
-    trunk = Lux.Chain(Lux.Dense(1, 32, Lux.σ), Lux.Dense(32, 32, Lux.σ), Lux.Dense(32, 1))
-    deeponet = NeuralPDE.DeepONet(branch, trunk)
+    bounds = (p = [0.1f0, pi],)
 
-    θ, st = Lux.setup(Random.default_rng(), deeponet)
-    a = rand(1, 10)
-    t = rand(1, 10)
-    x = (branch = a, trunk = t)
-    y, st = deeponet(x, θ, st)
-
-    bounds = (p = [0.1, pi / 2],)
-    #instances_size = 100 TODO remove dt -> instances_size
-    opt = OptimizationOptimisers.Adam(0.1)
-    alg = NeuralPDE.PINOODE(deeponet, opt, bounds)
-    sol, phi  = solve(prob, alg, dt = 0.1, verbose = true, maxiters = 2000)
-
-    instances_size = 100
-    p = range(bounds.p[1], stop = bounds.p[2], length = instances_size)
-    t = range(tspan[1], stop = tspan[2], length = instances_size)
-    x = (branch = collect(p)', trunk = collect(t)')
-    predict = phi(x, sol.u)
-
-    ground_func = (u0, p, t) -> u0 + sin(p * t) / (p)
-    ground_solution = ground_func.(u0, x.branch', x.trunk)
-    @test ground_solution ≈ predict atol=1
-end
+    strategy  = NeuralPDE.SomeStrategy(branch_size = 50, trunk_size = 40)
 
-@testset "Example with data" begin
-    equation = (u, p, t) -> cos(p * t)
-    tspan = (0.0f0, 2.0f0)
-    u0 = 0.0f0
-    prob = ODEProblem(linear, u0, tspan)
+    opt = OptimizationOptimisers.Adam(0.03)
+    alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy)
 
-    # init neural operator
-    deeponet = DeepONet(branch, trunk)
-    opt = OptimizationOptimisers.Adam(0.01)
-    bounds = (p = [0, pi / 2])
-    function data_loss()
-        #code
-    end
-    alg = NeuralPDE.PINOODE(chain, opt, bounds; add_loss = data_loss)
     sol = solve(prob, alg, verbose = false, maxiters = 2000)
-    predict = sol.predict
 
-    ground_func = (u0, p, t) -> u0 + sin(p * t) / (p)
-    ground = ground_func(..)
-    @test ground≈predict atol=1
-end
+    ground_analytic = (u0, p, t) ->  u0 + sin(p * t) / (p)
+    p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size)
+    p = reshape(p_, 1, branch_size, 1)
+    ground_solution = ground_analytic.(u0, p, sol.t.trunk)
 
-@testset "Example u0" begin
-    equation = (u, p, t) -> cos(p * t)
-    tspan = (0.0f0, 2.0f0)
-    p = Float32(pi)
-    # u0 = 2.0f0
-    prob = ODEProblem(linear, 0, tspan, p)
-    prob = PINOODEProblem(linear, tspan, p)
-
-    neuraloperator = DeepONet(branch, trunk)
-
-    opt = OptimizationOptimisers.Adam(0.001)
-    bounds = (u0 = [0, 2],)
-    alg = PINOODE(chain, opt, bounds)
-    pino_solution = solve(prob, alg, verbose = false, maxiters = 2000)
-    predict = pino_solution.predict
-
-    ground_func = (u0, p, t) -> u0 + sin(p * t) / (p)
-    ground = ground_func(...)
-    @test ground≈predict atol=1.0
+    @test ground_solution≈sol.u rtol=0.01
 end
 
-@testset "Example u0 and p" begin
-    equation = (u, p, t) -> cos(p * t)
-    tspan = (0.0f0, 2.0f0)
-    p = Float32(pi)
-    # u0 = 2.0f0
-    prob = ODEProblem(linear, 0, tspan, p)
-    prob = PINOODEProblem(linear, tspan, p)
-
-    neuraloperator = DeepONet(branch, trunk)
-
-    opt = OptimizationOptimisers.Adam(0.001)
-    bounds = (u0 = [0, 2],)
-    alg = PINOODE(chain, opt, bounds)
-    pino_solution = solve(prob, alg, verbose = false, maxiters = 2000)
-    predict = pino_solution.predict
-
-    ground_func = (u0, p, t) -> u0 + sin(p * t) / (p)
-    ground = ground_func(...)
-    @test ground≈predict atol=1.0
+@testset "Example du = p*t^2 " begin
+    equation = (u, p, t) -> p * t^2
+    tspan = (0.0f0, 1.0f0)
+    u0 = 0.f0
+    prob = ODEProblem(equation, u0, tspan)
+
+    # init neural operator
+    branch = Lux.Chain(
+        Lux.Dense(1, 10, Lux.tanh_fast),
+        Lux.Dense(10, 10, Lux.tanh_fast),
+        Lux.Dense(10, 10, Lux.tanh_fast))
+    trunk = Lux.Chain(
+        Lux.Dense(1, 10, Lux.tanh_fast),
+        Lux.Dense(10, 10, Lux.tanh_fast),
+        Lux.Dense(10, 10, Lux.tanh_fast))
+    linear = Lux.Chain(Lux.Dense(10, 1))
+    deeponet = NeuralPDE.DeepONet(branch, trunk; linear= linear)
+
+    a = rand(1, 50, 40)
+    b = rand(1, 1, 40)
+    x = (branch = a, trunk = b)
+    θ, st = Lux.setup(Random.default_rng(), deeponet)
+    c = deeponet(x, θ, st)[1]
+
+    bounds = (p = [0.0f0, 2.f0],)
+    strategy = NeuralPDE.SomeStrategy(branch_size = 50, trunk_size = 40)
+    opt = OptimizationOptimisers.Adam(0.03)
+    alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy)
+
+    sol = solve(prob, alg, verbose = false, maxiters = 2000)
+    ground_analytic = (u0, p, t) -> u0 + p * t^3 / 3
+    p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size)
+    p = reshape(p_, 1, branch_size, 1)
+    ground_solution = ground_analytic.(u0, p, sol.t.trunk)
+
+    @test ground_solution≈sol.u rtol=0.01
 end
+
+# @testset "Example with data" begin
+# end
diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl
deleted file mode 100644
index 919e6012ed..0000000000
--- a/test/PINO_ode_tests_gpu.jl
+++ /dev/null
@@ -1,151 +0,0 @@
-using Test
-using OrdinaryDiffEq
-using Lux
-using ComponentArrays
-#using NeuralOperators
-using OptimizationOptimisers
-using Random
-using LuxCUDA
-using NeuralPDE
-
-CUDA.allowscalar(false)
-const gpud = gpu_device()
-
-@testset "Example p" begin
-    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 / 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.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)
-        reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1))
-        push!(prob_set, prob_)
-        u_output_[:, :, i] = reshape_sol
-    end
-    u_output_ = u_output_ |> gpud
-
-    """
-    Set of 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_)
-
-    inner = 50
-    chain = Lux.Chain(Lux.Dense(2, inner, Lux.σ),
-        Lux.Dense(inner, inner, Lux.σ),
-        Lux.Dense(inner, inner, Lux.σ),
-        Lux.Dense(inner, inner, Lux.σ),
-        Lux.Dense(inner, inner, Lux.σ),
-        Lux.Dense(inner, 1))
-    ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud
-    opt = OptimizationOptimisers.Adam(0.03)
-    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
-    function lotka_volterra(u, p, t)
-        # Model parameters.
-        α, β, γ, δ = p
-        # Current state.
-        x, y = u
-        # Evaluate differential equations.
-        dx = (α - β * y) * x # prey
-        dy = (δ * x - γ) * y # predator
-        return [dx, dy]
-    end
-
-    u0 = [1.0f0, 1.0f0]
-    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 - 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)
-        reshape_sol_ = reduce(hcat, solution(range_).u)
-        reshape_sol = Float32.(reshape(reshape_sol_, 2, instances_size, 1))
-        push!(prob_set, prob_)
-        u_output_[:, :, i] = reshape_sol
-    end
-
-    train_set = TRAINSET(prob_set, u_output_)
-
-    # flat_no = FourierNeuralOperator(ch = (5, 64, 64, 64, 64, 64, 128, 2), modes = (16,),
-    #     σ = gelu)
-    # flat_no = Lux.transform(flat_no)
-    inner = 50
-    chain = Lux.Chain(Lux.Dense(5, inner, Lux.σ),
-        Lux.Dense(inner, inner, Lux.σ),
-        Lux.Dense(inner, inner, Lux.σ),
-        Lux.Dense(inner, inner, Lux.σ),
-        Lux.Dense(inner, inner, Lux.σ),
-        Lux.Dense(inner, 2))
-    ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud
-
-    opt = OptimizationOptimisers.Adam(0.001)
-    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
diff --git a/test/runtests.jl b/test/runtests.jl
index ac8a4b6e36..0990c63bb7 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -99,11 +99,6 @@ end
             include("NNPDE_tests_gpu_Lux.jl")
         end
     end
-    if !is_APPVEYOR && GROUP == "PINO_GPU"
-        @safetestset "PINO ode gpu" begin include("PINO_ode_tests_gpu.jl")
-        end
-    end
-
 
     if GROUP == "All" || GROUP == "DGM"
         @time @safetestset "Deep Galerkin solver" begin