如何在Flux.jl中定义自定义训练循环

using Flux
using Flux: @epochs, @epochs
using Flux: params, gradient

model = Chain(
    Dense(10, 32, relu),
    Dense(32, 2),
    softmax
)

loss(x, y) = Flux.crossentropy(model(x), y)

function train!(model, data, optimizer)
    for (x, y) in data
        gs = gradient(params(model)) do
            l = loss(x, y)
        end
        Flux.update!(optimizer, params(model), gs)
    end
end

data = DataLoader(train_data, batchsize=32, shuffle=true)

optimizer = ADAM()

@epochs 10 train!(model, data, optimizer)