TFSegformer fail to convert to onnx 
Author: joihn
Pull request: no
Labels: name: bug

Thanks for this repo  !

### Error Description

1) train a tensorflow segformer model 
2) save the checkpoint
2) try to convert it to onnx via
`python -m transformers.onnx --model CHECKPOINT_PATH OUT_PATH`
###  Error:

```
2022-08-22 12:46:58.488185: I tensorflow/stream_executor/cuda/cuda_dnn.cc:368] Loaded cuDNN version 8400
2022-08-22 12:47:00.216906: E tensorflow/stream_executor/cuda/cuda_blas.cc:197] failed to set new cublas math mode: CUBLAS_STATUS_INVALID_VALUE
2022-08-22 12:47:00.216958: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at matmul_op_impl.h:442 : INTERNAL: Failed initializing math mode
Traceback (most recent call last):
  File "/usr/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/usr/local/lib/python3.8/dist-packages/transformers/onnx/__main__.py", line 107, in <module>
    main()
  File "/usr/local/lib/python3.8/dist-packages/transformers/onnx/__main__.py", line 73, in main
    model = FeaturesManager.get_model_from_feature(
  File "/usr/local/lib/python3.8/dist-packages/transformers/onnx/features.py", line 579, in get_model_from_feature
    model = model_class.from_pretrained(model, from_tf=True, cache_dir=cache_dir)
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/auto/auto_factory.py", line 463, in from_pretrained
    return model_class.from_pretrained(
  File "/usr/local/lib/python3.8/dist-packages/transformers/modeling_utils.py", line 2194, in from_pretrained
    model, loading_info = load_tf2_checkpoint_in_pytorch_model(
  File "/usr/local/lib/python3.8/dist-packages/transformers/modeling_tf_pytorch_utils.py", line 343, in load_tf2_checkpoint_in_pytorch_model
    tf_model(tf_inputs, training=False)  # Make sure model is built
  File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/usr/local/lib/python3.8/dist-packages/transformers/modeling_tf_utils.py", line 407, in run_call_with_unpacked_inputs
    return func(self, **unpacked_inputs)
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/segformer/modeling_tf_segformer.py", line 630, in call
    outputs = self.segformer(
  File "/usr/local/lib/python3.8/dist-packages/transformers/modeling_tf_utils.py", line 407, in run_call_with_unpacked_inputs
    return func(self, **unpacked_inputs)
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/segformer/modeling_tf_segformer.py", line 488, in call
    encoder_outputs = self.encoder(
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/segformer/modeling_tf_segformer.py", line 422, in call
    layer_outputs = blk(
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/segformer/modeling_tf_segformer.py", line 324, in call
    self_attention_outputs = self.attention(
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/segformer/modeling_tf_segformer.py", line 232, in call
    self_outputs = self.self(hidden_states, height, width, output_attentions)
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/segformer/modeling_tf_segformer.py", line 161, in call
    query_layer = self.transpose_for_scores(self.query(hidden_states))
tensorflow.python.framework.errors_impl.InternalError: Exception encountered when calling layer "query" (type Dense).

Failed initializing math mode [Op:MatMul]

Call arguments received:
  • inputs=tf.Tensor(shape=(3, 16384, 32), dtype=float32)
```



### System Info
Ubuntu 20.04, Python 3.9, TF 2.9.1, Nvidia Titan

Who can help?
@sayakpaul @NielsRogge

Comments:
Not sure what's missing, but here's a notebook that shows the ONNX conversion (using `tf2onnx`): https://github.com/deep-diver/segformer-tf-transformers/blob/main/notebooks/TFSegFormer_ONNX.ipynb. 

I can confirm this approach is working, thanks for the quick answer