Выведите выходной тензор, и проверьте его тип данных, больше похоже на баг самого TensorFlow

Если действительно проблема в выходном тензоре попробуйте вот это преобразовании

val outputArray = when (rawOutputBuffer?.dataType) {
DataType.FLOAT32 -> rawOutputBuffer.floatArray.map { it.toInt() }.toIntArray() // Convert FloatArray to IntArray
else -> throw IllegalArgumentException("Unsupported output tensor data type: ${rawOutputBuffer?.dataType}")

}

Надеюсь вы решите свою проблему ))