switch to `TrainState` and/or `single_train_step!`

[lazarusA]

I think that replicates to some degree our steps, allowing us to call AutoEnzyme() or any AD package backend:

l, backtrace = Zygote.pullback((ps) -> lossfn(hybridModel, x, (y, is_no_nan), ps, st,
                        LoggingLoss(training_loss=training_loss, agg=agg)), ps)
grads = backtrace(l)[1]
Optimisers.update!(opt_state, ps, grads)
st =(; l[2].st...)

[Qfl3x]

Closed by #198