java - Error org.deeplearning4j.exception.DL4JInvalidInputException because nIn does not match sequence length - Stack Overflow

I want to teach a NN some time series with 2 features using Deeplearning4j in Kotlin. I use one LSTM Layer, the time series is 5 long and I picked a batch size of 16, so the shape of the training data is [16, 5, 2]. For the output I pick [1, 5, 2] since I am only interested in one single value. This is the code for initialization:

@Test
fun testRandomLearning() {
    // Create NN
    val batchSize = 16
    val sequenceLength = 5
    val numFeatures = 2

    val config = NeuralNetConfiguration.Builder()
        .seed(seed.toLong())
        .activation(Activation.RELU)
        .weightInit(org.deeplearning4j.nn.weights.WeightInit.VAR_SCALING_NORMAL_FAN_AVG)
        .updater(org.nd4j.linalg.learning.config.Adam(0.001))
        .list()
        .layer(0, LSTM.Builder()
            .nIn(numFeatures)
            .nOut(64)
            .activation(Activation.TANH)
            .build()
        )
        .layer(1, RnnOutputLayer.Builder()
            .nIn(64)
            .nOut(1)
            .activation(Activation.IDENTITY)
            .lossFunction(LossFunctions.LossFunction.MSE)
            .build()
        )
        .build()

    val model = MultiLayerNetwork(config).apply {
        init()
        setListeners(ScoreIterationListener(1000))
    }

    val numBatches = 10

    // Create random training data
    val dataList = (0 until numBatches).map {
        val features: INDArray = Nd4j.rand(*longArrayOf(batchSize.toLong(), sequenceLength.toLong(), numFeatures.toLong()))
        val labels: INDArray = Nd4j.rand(*longArrayOf(batchSize.toLong(), 1L))
        org.nd4jmon.primitives.Pair(features, labels)
    }

    val iterator: DataSetIterator = INDArrayDataSetIterator(dataList, batchSize)
    model.fit(iterator) // 2 epochs
    model.fit(iterator)

    // Check model
    val input = Nd4j.rand(*longArrayOf(1, sequenceLength.toLong(), numFeatures.toLong()))
    val output = model.output(input) // << == here
    println("$output")
}

However model.output(input) whose input matches that one of the training data except for batch size causes an exception that I don't understand:

org.deeplearning4j.exception.DL4JInvalidInputException: Received input with size(1) = 5 (input array shape = [1, 5, 2]); input.size(1) must match layer nIn size (nIn = 2)
at org.deeplearning4j.nn.layers.recurrent.LSTMHelpers.activateHelper(LSTMHelpers.java:171)

input.size(1) is the sequence length, and according to this message it should match the value passed to nIn in layer, but nIn should be the number of features which is clearly 2. Do I have some deeper misunderstanding here?

Used versions:

implementation("org.deeplearning4j:deeplearning4j-core:1.0.0-M2.1")
implementation("org.nd4j:nd4j-native-platform:1.0.0-M2.1")