You can see what the model is predicting for the test data with Axon.predict/4
Axon.predict(model, trained_model_state, x_test)
You can see what the model is predicting for the test data with Axon.predict/4
Axon.predict(model, trained_model_state, x_test)