The model training loop on page 273 should visualize the model output after each epoch. Instead, an argument error is reported after the first epoch:
** (ArgumentError) expected Nx.Tensor to have shape {height, width, channels}, got: {1, 28, 28, 1}
(kino 0.12.3) lib/kino/image.ex:88: Kino.Image.new/1
#cell:vrknsgv5b46le6ni:16: (file)
#cell:vrknsgv5b46le6ni:5: (file)
The problem arises in the visualize_test_image
function spanning pages 271-272. Axon.predict
returns an image tensor with dimensions {1, 28, 28, 1}
. This tensor can be reshaped to the {height, width, channels}
dimensions needed for Kino.Image.new
by changing the line at the top of page 272 to:
out_image = Nx.multiply(out_image, 255) |> Nx.as_type(:u8) |> Nx.reshape({28, 28, 1})
With this change, the visualizations appear as shown in the book on page 273.