kMaX onnx export and trt export fails

[hannes09]

Hi,

I am trying to export a kMaX model to onnx format or a trt engine.
The error accures when calling trt.TrtGraphConverterV2
Both methods fail with the same error:

StatefulPartitionedCall/kMaXDeepLab/kmax_pixel_decoder/kmax_pixeldecoder_0/block1/attention/height_axis/query_rpe/Gather/axis is not unique

I am using kmax_meta_r50_os32.textproto with a dataset similar to cityscape.

The same error was mentioned in https://github.com/google-research/deeplab2/issues/50
and in https://github.com/google-research/deeplab2/issues/82

I have read the answers but for me it is stillnot clear how to fix the issue.

I only need support a batch size of 1.