When training the cnn_model, I get the following output:

Epoch: 0, Batch: 150, accuracy: 0.4985513 loss: 7.6424022
Epoch: 1, Batch: 163, accuracy: 0.4992854 loss: 7.6783161
Epoch: 2, Batch: 176, accuracy: 0.5000441 loss: 7.6865749
Epoch: 3, Batch: 139, accuracy: 0.4983259 loss: 7.6991839
Epoch: 4, Batch: 152, accuracy: 0.4988766 loss: 7.6995916

%{
  "conv_0" => %{
    "bias" => #Nx.Tensor<
      f32[32]
      EXLA.Backend<host:0, 0.1357844422.1979580433.82179>
      [NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN]
    >,
    "kernel" => #Nx.Tensor<
      f32[3][3][3][32]
      EXLA.Backend<host:0, 0.1357844422.1979580433.82180>
      [
        [
          [
            [NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN],
            [NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, ...],
            ...
          ],
          ...
        ],
        ...
      ]
    >
  },
  "conv_1" => %{
    "bias" => #Nx.Tensor<
      f32[64]
      EXLA.Backend<host:0, 0.1357844422.1979580433.82181>
      [NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, -0.0071477023884654045, NaN, NaN, NaN, NaN, 0.0, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, ...]
    >,
    "kernel" => #Nx.Tensor<
      f32[3][3][32][64]
      EXLA.Backend<host:0, 0.1357844422.1979580433.82182>
      [
        [
          [
            [NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, ...],
            ...
          ],
          ...
        ],
        ...
      ]
    >
  },
  "conv_2" => %{
    "bias" => #Nx.Tensor<
      f32[128]
      EXLA.Backend<host:0, 0.1357844422.1979580433.82183>
      [0.0, NaN, NaN, NaN, NaN, NaN, NaN, NaN, 0.005036031361669302, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, ...]
    >,
    "kernel" => #Nx.Tensor<
      f32[3][3][64][128]
      EXLA.Backend<host:0, 0.1357844422.1979580433.82184>
      [
        [
          [
            [NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, ...],
            ...
          ],
          ...
        ],
        ...
      ]
    >
  },
  "dense_0" => %{
    "bias" => #Nx.Tensor<
      f32[128]
      EXLA.Backend<host:0, 0.1357844422.1979580433.82185>
      [NaN, -0.005992305930703878, -0.006005365401506424, -0.004664595704525709, NaN, NaN, NaN, -5.619042203761637e-4, 0.0, NaN, -0.005999671295285225, -6.131592726887902e-6, NaN, 0.0, NaN, 0.0, 0.0, NaN, NaN, -0.006002828478813171, -0.00600335793569684, 0.0, NaN, NaN, NaN, -0.006002923008054495, -0.006005282513797283, -0.00600528996437788, -0.0060048955492675304, -0.006004981696605682, NaN, -0.006004655733704567, -0.006005233619362116, NaN, -0.006004724185913801, -0.006005335133522749, -0.006005051080137491, -0.006004408933222294, NaN, -0.006005355156958103, 0.0, -0.006005344912409782, 0.0, NaN, -0.005991040728986263, ...]
    >,
    "kernel" => #Nx.Tensor<
      f32[18432][128]
      EXLA.Backend<host:0, 0.1357844422.1979580433.82186>
      [
        [NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, ...],
        ...
      ]
    >
  },
  "dense_1" => %{
    "bias" => #Nx.Tensor<
      f32[1]
      EXLA.Backend<host:0, 0.1357844422.1979580433.82187>
      [NaN]
    >,
    "kernel" => #Nx.Tensor<
      f32[128][1]
      EXLA.Backend<host:0, 0.1357844422.1979580433.82188>
      [
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        [NaN],
        ...
      ]
    >
  }
}

The accuracy of the mlp_model was Batch: 6, accuracy: 0.5078125 and the accuracy of this cnn_model is Batch: 6, accuracy: 0.4944196 which was slightly worse instead of the expected “significantly better.”

I reviewed all the code to make sure I hadn’t missed anything, but I couldn’t find anything that didn’t match.

I’m guessing the NaNs in the trained model state are a problem, but I’m not sure how to fix that.