Simple neural network in Julia

Code

using Flux, MLDatasets, CUDA, FileIO
using Flux: train!, onehotbatch
x_train, y_train = MLDatasets.MNIST.traindata()
x_test, y_test = MLDatasets.MNIST.testdata()

x_train = Float32.(x_train)
y_train = Flux.onehotbatch(y_train, 0:9)
model = Chain(
    Dense(784, 256, relu),
    Dense(256, 10, relu), softmax
)

loss(x, y) = Flux.Losses.logitcrossentropy(model(x), y)
optimizer = ADAM(0.0001)
parameters = params(model)
train_data = [(Flux.flatten(x_train), y_train)]
test_data = [(Flux.flatten(x_test), y_test)]

for i in 1:400
    Flux.train!(loss, parameters, train_data, optimizer)
end

accuracy = 0
for i in 1:length(y_test)
    if findmax(model(test_data[1][1][:, i]))[2] - 1  == y_test[i]
        accuracy = accuracy + 1
    end
end

println(accuracy / length(y_test))

Hi!

I'm trying to execute this exact code, but I get the following error:

ERROR: LoadError: MethodError: no method matching _methods_by_ftype(::Type{Tuple{typeof(ChainRulesCore.rrule),Flux.Optimise.var"#15#21"{typeof(loss),Tuple{Array{Float32,2},Flux.OneHotMatrix{Array{Flux.OneHotVector,1}}}}}}, ::Int64, ::UInt64, ::Bool, ::Base.RefValue{UInt64}, ::Base.RefValue{UInt64}, ::Ptr{Int32})
Closest candidates are:
_methods_by_ftype(::Any, ::Int64, ::UInt64) at reflection.jl:838
_methods_by_ftype(::Any, ::Int64, ::UInt64, !Matched::Array{UInt64,1}, !Matched::Array{UInt64,1}) at reflection.jl:841

Could someone tell me why please?

Thank you in advance :)