Update src/pino_ode_solve.jl

Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com>
SciML · Jul 17, 2024 · 8ae7b0f · 8ae7b0f
1 parent 09e5555
commit 8ae7b0f
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl
@@ -67,11 +67,8 @@ end
 
 function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params)
  θ, st = Lux.setup(Random.default_rng(), chain)
- if init_params === nothing
- init_params = ComponentArrays.ComponentArray(θ)
- else
- init_params = ComponentArrays.ComponentArray(init_params)
- end
+ init_params = isnothing(init_params) ? θ : init_params
+ init_params = ComponentArrays.ComponentArray(init_params)
  PINOPhi(chain, st), init_params
 end