Difference in Inference result in JAVA and Swift using same TFLite Model

I have trained a Mobile Bert Custom TFLite model for Text Classification and the model works as expected in Android. The output values match the expected values from running the TFLite model in Python. However, while running the same model in iOS with same input, I am getting different results.

I am new to iOS, this is the first time I’m working on iOS so my gut feeling is that I’m doing something wrong. The only clue I currently have is the slight difference in my Java Code and Swift Code to extract the result from Model Output.

Java

float[][][] bert_output = new float[1][MAX_SEQ_LEN][BERT_HIDDEN_LAYERS]; 
bert_model.run(inputIds, bert_output);

inputIds here is an Int Array
In Java I directly call run in the interpreter and pass the Input in inputIds as well as the output array object and at the end of it I get a populated Output array in bert_output.

Swift

In Swift the inference seems to work a little different. You have to convert Your input(Int Array) to Tensor format and also fetch the output in Tensor format and convert it back to Float array. I feel this is the part where things are going wrong but no way to confirm.

Convert Input into Tensor

let input_tensor = inputIds.withUnsafeBufferPointer(Data.init)

Copy Input to Model Input Layer 0

try interpreter_bert?.copy(input, toInputAt: 0)

Inference

try interpreter_bert?.invoke()

Copy Output Tensor from Output Layer 0

let temp_res = try interpreter_bert?.output(at: 0)

Convert Output Tensor Data to Float Array

let mResult = BertHelper.TensorToArray<Float32>(tensor: temp_res!)

You can find the Conversion code here

After the Conversion, the Result Shape matches the Java result and Expected result but the probability values don’t match up.

It would be a big help if someone here could help me. Tensorflow documentation and examples where Model Interpreter is directly called, only showcase Models which work on images and the NLP examples use TensorflowTaskLibrary, hence putting my issue up here.