如何在Flux.jl中定义自定义损失函数？

using Flux

# 定义自定义损失函数
function custom_loss(y_pred, y_true)
    # 计算损失
    loss = sum(abs2, y_pred - y_true)
    return loss
end

# 创建模型
model = Chain(Dense(10, 5, relu), Dense(5, 1))

# 定义损失函数
loss(x, y) = custom_loss(model(x), y)

# 生成训练数据
x_train = rand(10, 10)
y_train = rand(10)

# 训练模型
Flux.train!(loss, params(model), [(x_train, y_train)], ADAM())

# 使用模型进行预测
x_test = rand(5, 10)
y_pred = model(x_test)