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]






















