# Specify the model
model = Flux.Chain(
    Flux.Dense(size(xtrn, 2), 10, Flux.relu),
    Flux.Dense(10, 3),
    Flux.softmax
);

# Define loss function
loss(x, y) = Flux.crossentropy(model(x'), Float32.(y));