[Book Question] Machine Learning in Elixir: Poor accuracy for Chapter 1's example

I had a few more minutes to play with this, and so far it looks like we can get better results by processing the x and y data into tensors before splitting up test and training sets.

feature_columns = ["sepal_length", "sepal_width", "petal_length", "petal_width"]
label_column = "species"

x_all = Nx.stack(shuffled_normalized_iris[feature_columns], axis: 1)

y_all =
  shuffled_normalized_iris[label_column]
  |> Explorer.Series.cast(:category)
  |> Nx.stack(axis: -1)
  |> Nx.equal(Nx.iota({1, 3}, axis: -1))

x_train = x_all[0..119]
x_test = x_all[120..149]

y_train = y_all[0..119]
y_test = y_all[120..149]