But that is not what the cost function expects, which is the point of those comments in the assignment. So how would one deal with that? Have you considered applying the “transpose” operation to the labels and logits before passing them to the cost function?
Also note that the Z3 outputs are “logits” and not softmax output values. So you have two choices 1) apply softmax yourself or 2) use the from_logits argument to tell the cost function to do that computation internally. Option 2) is the better choice, which is why the given code does not include the softmax logic.