fixup! refactor: PINOODEInterpolation

SciML · Feb 3, 2025 · 2dc09cf · 2dc09cf
1 parent 2642a92
commit 2dc09cf
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl
@@ -29,7 +29,7 @@ using Reexport: @reexport
 using RuntimeGeneratedFunctions: RuntimeGeneratedFunctions, @RuntimeGeneratedFunction
 using SciMLBase: SciMLBase, BatchIntegralFunction, IntegralProblem, NoiseProblem,
                  OptimizationFunction, OptimizationProblem, ReturnCode, discretize,
-                 isinplace, solve, symbolic_discretize, ODEProblem
+                 isinplace, solve, symbolic_discretize, ODEProblem, AbstractODESolution
 using Statistics: Statistics, mean
 using QuasiMonteCarlo: QuasiMonteCarlo, LatinHypercubeSample
 using WeightInitializers: glorot_uniform, zeros32

diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl
@@ -292,7 +292,7 @@ end
 SciMLBase.interp_summary(::PINOODEInterpolation) = "Trained neural network interpolation"
 SciMLBase.allowscomplex(::PINOODE) = true
 
-function (sol::SciMLBase.AbstractODESolution)(
+function (sol::AbstractODESolution)(
         t::AbstractArray, ::Type{deriv}, idxs::Nothing,
         continuity) where {deriv}
     sol.interp(t, idxs, deriv, sol.prob.p, continuity)