The intent of the feedForwardWithKey() method in step 1 and 2 is to generate output/predictions for the given input dataset. A map is returned from this method. The input data is represented by the keys in the map and the results (output) are represented by values (INDArray).

feedForwardWithKey() accepts two attributes: input data and the minibatch size for feed-forward operations. The input data (features) is in the format of JavaPairRDD<K, INDArray>.

Note that RDD data is unordered. We need a way to map each input to the respective results (output). Hence, we need to have a key-value pair that maps each input to its respective output. That's the main reason why we use key values here. It has nothing to do with the inference process. Values for the minibatch size are used for the trade-off between memory versus computational efficiency.