Axon ignoring the activation function? about axon HOT 5 CLOSED

train_data = Scidata.MNIST.download()
test_data = Scidata.MNIST.download_test()

{train_data, train_labels} = train_data
{train_binary, train_type, train_shape} = train_data
{train_label_binary, train_label_type, train_label_shape} = train_labels

train_y =
  Nx.from_binary(train_label_binary, train_label_type)
  |> Nx.reshape(train_label_shape)
  |> Nx.new_axis(-1)
  |> Nx.equal(Enum.to_list(0..9) |> Nx.tensor())
  |> Nx.to_batched(32)
  |> Enum.to_list()

train_x =
  Nx.from_binary(train_binary, train_type)
  |> Nx.reshape(train_shape)
  |> Nx.divide(255)
  |> Nx.to_batched(32)
  |> Enum.to_list()

data = Enum.zip(train_x, train_y)

training_count = floor(0.8 * Enum.count(data))
validation_count = floor(0.2 * training_count)

{training_data, test_data} = Enum.split(data, training_count)
{validation_data, training_data} = Enum.split(training_data, validation_count)

model =
  Axon.input("input_0", shape: {nil, 1, 28, 28})
  |> Axon.flatten()
  |> Axon.dense(128, activation: :relu)
  |> Axon.dense(10, activation: :sigmoid)

state =
  model
  |> Axon.Loop.trainer(:categorical_cross_entropy, Axon.Optimizers.adam(0.01))
  |> Axon.Loop.metric(:accuracy, "Accuracy")
  |> Axon.Loop.validate(model, validation_data)
  |> Axon.Loop.run(training_data, %{}, compiler: EXLA, epochs: 10)

model
|> Axon.Loop.evaluator()
|> Axon.Loop.metric(:accuracy, "Accuracy")
|> Axon.Loop.run(test_data, state)

first_number = Enum.at(train_x, 0)[0]
Axon.predict(model, state, first_number)

#Nx.Tensor<
  f32[1][10]
  EXLA.Backend<host:0, 0.3007448411.1655570452.243887>
  [
    [2.3431260944184697e-23, 4.296128036651581e-11, 1.0659236316176752e-17, 0.9999992847442627, 6.712944780263512e-22, 1.0, 2.8145804840526482e-22, 2.4531429665408666e-10, 9.463095196338145e-9, 2.6594868813845096e-6]
  ]
>