There is an issue with how I transcribed one-hot encoding in the book. The correct version:
train_df
|> DF.pull(label_column)
|> Explorer.Series.to_list()
|> Enum.map(fn
"Iris-setosa" -> 0
"Iris-versicolor" -> 1
"Iris-virginica" -> 2
end)
|> Nx.tensor(type: :u8)
|> Nx.new_axis(-1)
|> Nx.equal(Nx.iota({1, 3}, axis: -1))