Hello @sandra_jayan,
Notice please the shape of the arguments logits and labels as received by the compute_cost function and the shape expected by tf.keras.losses.categorical_crossentropy, maybe you need to adjust that. Also be aware of the from_logits parameter needed in the call to the categorical_crossentropy function.