Semantic Segementation Example "Label number 21 mismatch" exception

When Following the example form this page :
www.tensorflow.org/lite/examples/segmentation/overview
I downloaded the linked .tflite model LiteModelDeeplabv31Metadata2 and imported it into my Android Studio project
using the right-click -> new -> other -> Tensor Flow Lite Model option.
The IDE Recognized the model and gave me the following code example to use it :

LiteModelDeeplabv31Metadata2 model = LiteModelDeeplabv31Metadata2.newInstance(context);

// Creates inputs for reference.
TensorImage image = TensorImage.fromBitmap(bitmap);

// Runs model inference and gets result.
LiteModelDeeplabv31Metadata2.Outputs outputs = model.process(image);
List<Category> segmentationMasks = outputs.getSegmentationMasksAsCategoryList(); 

When trying to execute this i get the following exception:

    java.lang.IllegalArgumentException: Label number 21 mismatch the shape on axis 1
        at org.tensorflow.lite.support.common.internal.SupportPreconditions.checkArgument(SupportPreconditions.java:104)
        at org.tensorflow.lite.support.label.TensorLabel.<init>(TensorLabel.java:87)
        at org.tensorflow.lite.support.label.TensorLabel.<init>(TensorLabel.java:105)
        at de.tudortmund.mb.ips.ml.LiteModelDeeplabv31Metadata2$Outputs.getSegmentationMasksAsCategoryList(LiteModelDeeplabv31Metadata2.java:108)

Looking through the stacktrace it becomes obvious what is going wrong :
getSegmentationMasksAsCategoryList(LiteModelDeeplabv31Metadata2.java:108)
is from the automatically generated code and
it creates a new new TensorLabel and calls .getCategoryList() on it.
However the constructor already fails:

This is the called Constructor :

public TensorLabel(@NonNull List<String> axisLabels, @NonNull TensorBuffer tensorBuffer) {
    this(makeMap(getFirstAxisWithSizeGreaterThanOne(tensorBuffer), axisLabels), tensorBuffer);
}

It receives a List of Labels as strings form the included meta data (21 strings),
creates a hashmap from it using the makeMap method with getFirstAxisWithSizeGreaterThanOne(tensorBuffer)
as its only key and the label list as the it's value.
This is strange however as the label dimension is not the FirstAxisWithSizeGreaterThanOne
but the last axis as per documentation.
The shape of the tensorBuffer is (1,257,257,21) where 257*257 is the resolution of the 21 semantic masks
which is the correct format according to this site:
www.tensorflow.org/lite/inference_with_metadata/task_library/image_segmenter

This wrongly created hashmap causes the actual constructor to fail in the last checkArgument call :

public TensorLabel(
      @NonNull Map<Integer, List<String>> axisLabels, @NonNull TensorBuffer tensorBuffer) {
    SupportPreconditions.checkNotNull(axisLabels, "Axis labels cannot be null.");
    SupportPreconditions.checkNotNull(tensorBuffer, "Tensor Buffer cannot be null.");
    this.axisLabels = axisLabels;
    this.tensorBuffer = tensorBuffer;
    this.shape = tensorBuffer.getShape();
    for (Map.Entry<Integer, List<String>> entry : axisLabels.entrySet()) {
      int axis = entry.getKey();
      SupportPreconditions.checkArgument(
          axis >= 0 && axis < shape.length, "Invalid axis id: " + axis);
      SupportPreconditions.checkNotNull(entry.getValue(), "Label list is null on axis " + axis);
      SupportPreconditions.checkArgument(
          shape[axis] == entry.getValue().size(),
          "Label number " + entry.getValue().size() + " mismatch the shape on axis " + axis);
    }
}

Either the called TensorLabel constructor has to be changed or the hashmap has to be created in advance and then the right
constructor can be called immediately (the one receiving a Map not a List)

Start by reading TensorLabel.java, including the List and Map constructors, then inspect the generated LiteModelDeeplabv31Metadata2.java entry point named in the stack trace. Reproduce the semantic-segmentation example with the documented (1,257,257,21) tensor and confirm the label axis is handled without the mismatch exception.