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 TFSegformerForSemanticSegmentation
model_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!