Problems running on GPU

[juliagpu]

Problems running on GPU

[jacobcvt12]

@carlobaldassi I had this same issue. It appears to be due to the numerical instability of calling softmax followed by crossentropy. Instead, using the function logitcrossentropy in your loss. As @boathit describes, this is equivalent but more stable. You will need to edit the code to something like

m = Chain(
  Dense(28^2, 32, relu),
  Dense(32, 10)) |> gpu

loss(x, y) = logitcrossentropy(m(x), y)

See FluxML/Flux.jl#201 as well.

[maleadt]

FluxML/Flux.jl#201 (comment)