Java TensorFlow Tensor Access

How can I find the index of the Max Float value in a Tensor using TensorFlow Java 0.2.0?

I am pretty new to Tensorflow Java and not yet clear on the relationship between NdArray and Tensor. If there is a better way to convert my List to the Tensor input, please let me know!

From SavedModelBundle I have retrieved a FLOAT (1) tensor with shape [1, 3] using the SessionRunner and feeding it with input tensors.

Wanted Behavior is as follows

text = "Love it"

val = Tensor of [0.12, 0.2, 0.68]

.predict(text) return --> 2

JAVA

public Integer predict(String text) {

        // Returns tokenized text in a List of MAXLEN
        List<Integer> token_ids = tokenize(text);

        // Convert List to tensor for Input
        IntDataBuffer bufferTokens = DataBuffers.ofInts(MAXLEN);
        int[] primArr = new int[MAXLEN];
        for (int i=0; i<MAXLEN; i++) {
            primArr[i] = token_ids.get(i);
        }
        bufferTokens.write(primArr);

        IntNdArray tokensMatrix = NdArrays.ofInts(Shape.of(1, MAXLEN));
        IntNdArray vector = tokensMatrix.get(0);
        vector.write(bufferTokens);

        Tensor<TInt32> input = TInt32.tensorOf(tokensMatrix);

        // Model.predict
        Tensor output = model.session()
                .runner()
                .feed("serving_default_input_ids:0", input)
                .fetch("StatefulPartitionedCall:0")
                .run() // List<Tensor<?>>
                .get(0);

        // TODO - HELP NEEDED: Extract arg max from tensor
        Tensor val = output.expect(TFloat32.DTYPE); // val = FLOAT (1) tensor with shape [1, 3]
        Integer maxIndex = ????????

        return maxIndex;
    }

The model Input, Output info when running myModel.metaGraphDef().getSignatureDefMap().get("serving_default"); is the following.

 ModelInfo
    inputs {
        key: "input_ids"
        value {
            name: "serving_default_input_ids:0"
            dtype: DT_INT32
            tensor_shape {
                dim {
                    size: -1
                }
                dim {
                    size: MAXLEN
                }
            }
        }
    }
    outputs {
        key: "dense_3"
        value {
            name: "StatefulPartitionedCall:0"
            dtype: DT_FLOAT
            tensor_shape {
                dim {
                    size: -1
                }
                dim {
                    size: 3
                }
            }
        }
    }

Thanks in advance!