Reduced accuracy with improved softmax implementation

[shubby007]

I was gettiing an accuracy of 95% with general implementation of softmax function when performing the MNIST handwritten digit classification but getting an accuracy of 45% with the improved implementation of softmax function.
What am I doing wrong in this?
I have attached jpeg files of both my codes.

Thanks!

[rmwkwok]

Please check if your metric accepts logits, or use a metric that accepts logits.

Raymond