Error in PredictCost() for the op: "Softmax"

[xingdi1990]

I am running this code on my 3090 with tensorflow-gpu 2.5.0 and cuda 11.2

I can successfully run the training, but I get the error as follows:

Can someone give me a hint to solve this?