*{moderator edit - solution code removed}*

I am not sure where I went wrong, not able to get the expected output. I took the transpose of the input, set from_logits = True to apply softmax activation. It is mentioned that to divide the sum by the total number of samples to get the cost value. But as it is used to calculate batch wise am not sure where to get the batch size to compute the average.