Can't save the tensorflow model of nvidia/mit-b5
When I try to save the tensorflow model of nvidia/mit-b5 I get a very cryptic error, related to a failed transpose (permuting 4 dimensions but given a single number).The model trains just fine, and also inference works after the model is trained. But it cannot save as a tensorflow saved_model. Note that it does save to H5 successfully (i.e. saved_model=False).The reproduction is very easy - simply creating it and trying to save it brings this error (my original code has specific classes that I fine tune to, but the same error appears with any creation of the model). Here’s a snippet that fails:from transformers import TFSegformerForSemanticSegmentationmodel_checkpoint = “nvidia/mit-b5”model = TFSegformerForSemanticSegmentation.from_pretrained(model_checkpoint,)model.save_pretrained(‘/workspace/saved_model’, saved_model=True)The error is the following:File /usr/local/lib/python3.11/dist-packages/transformers/models/segformer/modeling_tf_segformer.py:878, in TFSegformerDecodeHead.call(self, encoder_hidden_states, training)875     encoder_hidden_state = tf.reshape(encoder_hidden_state, (-1, height, width, channel_dim))877 # unify channel dimension→ 878 encoder_hidden_state = tf.transpose(encoder_hidden_state, perm=[0, 2, 3, 1])879 height, width = shape_list(encoder_hidden_state)[1:3]880 encoder_hidden_state = mlp(encoder_hidden_state)ValueError: Dimension must be 0 but is 4 for ‘{{node transpose}} = Transpose[T=DT_STRING, Tperm=DT_INT32](transpose/a, transpose/perm)’ with input shapes:, [4].Can anyone help me me find a workaround or can help me uncover the root cause?I am using tensorflow 2.16.1.Thanks in advance!