That basic method works, but you will need to add the softmax processing of the output layer. Note that this network was defined to produce the linear activation outputs at the output layer, so that it can use the from_logits = True mode of the loss function. Here’s a thread which talks about that.
2 Likes