diffeqflux to torchdyn : training block

[shivance]

Hi Folks, I'm currently working on a notebook tutorial on neural odes. I'm following a julia tutorial and porting it to pytorch as not much material for neural ode is available in python.

For my use case I need to port following block of julia code into Python, and I seek help regarding the same

  using OrdinaryDiffEq, Flux, Random

  function train_one_round(node, θ, y, opt, maxiters,
                           y0 = y[:, 1]; kwargs...)
      predict(θ) = Array(node(y0, θ))
      loss(θ) = begin
          ŷ = predict(θ)
          Flux.mse(ŷ, y)
      end

      θ = θ == nothing ? node.p : θ
      res = DiffEqFlux.sciml_train(
          loss, θ, opt,
          maxiters = maxiters;
          kwargs...
      )
      return res.minimizer
  end

  function train(θ = nothing, maxiters = 150, lr = 1e-2)
      log_results(θs, losses) =
          (θ, loss) -> begin
              push!(θs, copy(θ))
              push!(losses, loss)
              false
          end

      θs, losses = [], []
      num_obs = 4:4:length(train_t)
      for k in num_obs
          node = neural_ode(train_t[1:k], size(y, 1))
          θ = train_one_round(
              node, θ, train_y[:, 1:k],
              ADAMW(lr), maxiters;
              cb = log_results(θs, losses)
          )
      end
      θs, losses
  end

  Random.seed!(1)
  θs, losses = train();

Will be highly oblidged .