I need to transpose a 3-dimensional tensor of the shape (batch_size, N, M) to (batch_size, M, N) in a custom loss function in Keras with tensorflow as the backend. I tried using the following function


It gives me an error that the dimension must be 2 instead, 3 is given.

When I try to input the following


it says the dimension is 3, instead, 2 is given.

I also tried using K.transpose but it gives me (M,N,batch_size).

I already implemented the solution with for loop, but in keras custom loss, for loops should be avoided.

from keras import backend as K
def custom_loss(y_true, y_pred):
    beta  = K.permute_dimensions(y_pred, (0,2,1))
  • $\begingroup$ When you do y_true.shape what values do you get? Also give the values for each batch_size, N, M. $\endgroup$
    – Memristor
    Jun 14 at 15:37
  • $\begingroup$ @Memristor For y_true.shape I got (batch_size, N, M). $\endgroup$
    – Doc Jazzy
    Jun 14 at 15:44

1 Answer 1


Use tf.transpose instead:

tf.transpose(y_true, perm=(0, 2, 1))

Otherwise use tf.keras instead of simply keras, as follows:

from tensorflow.keras import backend as K

def custom_loss(y_true, y_pred):
    beta = K.permute_dimensions(y_pred, (0, 2, 1))
    # ...
  • $\begingroup$ both approaches didn't work out for me. Did you try it, if yes, please mention the TF and Keras version. $\endgroup$
    – Doc Jazzy
    Jun 14 at 18:15
  • $\begingroup$ @DocJazzy Yes I tried both on TF and Keras v2.8.0. What happens if you reshape the tensor? For example x = tf.reshape(x, shape=(M, N, batch_size)). $\endgroup$ Jun 15 at 10:19
  • $\begingroup$ Thank you for your reply! It was a problem with the keras version, I installed the latest version, and it works without any problem. $\endgroup$
    – Doc Jazzy
    Jun 16 at 15:54

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.