How to use a custom loss function with reactant?

nevermind sorry, a loss should be define like this

function loss_fn(model,ps,st,d)
    x,y = d
    ŷ,stn = model(x,ps,st)
    return Lux.MSELoss()(ŷ,y),stn,(;)
end
2 Likes