Machine Learning in Elixir: Error running last example in Chapter 1 (Evaluating the Trained Model)

@seanmor5

I am enjoying the book so far, but ran into an Elixir error on the last example in chapter 1:

In my last Livebook cell I have:

data = [{x_test, y_test}]

model
|> Axon.Loop.evaluator()
|> Axon.Loop.metric(:accuracy)
|> Axon.Loop.run(data, trained_model_state)```

and I get this error:

** (ArgumentError) cannot broadcast tensor of dimensions {120} to {30}
(nx 0.7.2) lib/nx/shape.ex:345: Nx.Shape.binary_broadcast/4
(nx 0.7.2) lib/nx.ex:5427: Nx.devectorized_element_wise_pred_op/3
(axon 0.6.1) lib/axon/metrics.ex:88: Axon.Metrics.accuracy_transform/4
(axon 0.6.1) lib/axon/metrics.ex:467: anonymous fn/4 in Axon.Metrics.running_average/1
(axon 0.6.1) lib/axon/loop.ex:1974: anonymous fn/5 in Axon.Loop.build_batch_fn/2
(elixir 1.15.7) lib/enum.ex:3945: anonymous fn/4 in Enum.zip_with/3
(elixir 1.15.7) lib/enum.ex:4005: anonymous fn/3 in Enum.zip_reduce/4
#cell:vs5uswexiozx5vnr:6: (file)

I typed in all the examples verbatim up to that point. Any ideas what is going on here?

Thanks,

Greg