TFLite android Error - Cannot copy from a TensorFlowLite tensor (StatefulPartitionedCall:2) with 131072 bytes to a Java Buffer with 2048 bytes

I’m new to Android and TF and I’m having an issue that I couldn’t find any solution for. I have a model in python that I converted to TFLite and imported into Android Studio. When I imported it, here’s the sample code suggested for me by Android Studio:

val model = Model.newInstance(context)

// Creates inputs for reference.
val inputFeature0 = TensorBuffer.createFixedSize(intArrayOf(1, 64, 112, 112, 3), DataType.FLOAT32)
inputFeature0.loadBuffer(byteBuffer)

// Runs model inference and gets result.
val outputs = model.process(inputFeature0)
val outputFeature0 = outputs.outputFeature0AsTensorBuffer
val outputFeature1 = outputs.outputFeature1AsTensorBuffer
val outputFeature2 = outputs.outputFeature2AsTensorBuffer

// Releases model resources if no longer used.
model.close()

I’m using the CameraX API to get the camera images and here is my analyze function:

        override fun analyze(imageProxy: ImageProxy) {

            val bitmap = imageProxy.toBitmap()
            val scaledBitmap = Bitmap.createScaledBitmap(bitmap, 112, 112, true)
            scaledBitmap.copyPixelsToBuffer(byteBuffer)
            inputFeature0.loadBuffer(byteBuffer)
            val outputs = model.process(inputFeature0)
            imageProxy.close()
        }

where byteBuffer and inputFeature0 are defined as members of my class like this:

        val byteBuffer = ByteBuffer.allocateDirect(1*64*112*112*3*4)
        val inputFeature0 = TensorBuffer.createFixedSize(intArrayOf(1, 64, 112, 112, 3), DataType.FLOAT32)

When the model.process is called I’m getting this error:

java.lang.IllegalArgumentException: Cannot copy from a TensorFlowLite tensor (StatefulPartitionedCall:2) with 131072 bytes to a Java Buffer with 2048 bytes.

131072 is 64 times the value of 2048 and 64 is the number of frames my model expects. I tried to change the byteBuffer size to 1121123*4 but I got this error:

java.lang.IllegalArgumentException: The size of byte buffer and the shape do not match. Expected: 9633792 Actual: 150528

which suggests that I had the size right before.

Any ideas?

Thanks