Does tf.keras.losses.categorical_crossentropy return an array or a single value?
Most usual losses return the original shape minus the last axis.
So, if your original y_pred shape was (samples, ..., ..., classes), then your resulting shape will be (samples, ..., ...).
This is probably because Keras may use this tensor in further calculations, for sample weights and maybe other things. 
In a custom loop, if these dimensions are useless, you can simply take a K.mean(loss_result) before calculating the gradients. (Where K is either keras.backend or tensorflow.keras.backend)
I'm using a custom training loop. The loss that is returned by tf.keras.losses.categorical_crossentropy is an array of I'm assuming (1,batch_size). Is this what it is supposed to return or a single value?
In the latter case, any idea what I could be doing wrong?