docs: public <K> JavaPairRDD<K, INDArray> feedForwardWithMaskAndKey(JavaPairRDD<K, Tuple2<INDArray,INDArray>> featuresDataAndMask, int batchSize) {
          return featuresDataAndMask
                  .mapPartitionsToPair(new FeedForwardWithKeyFunction<K>(sc.broadcast(network.params()),
                          sc.broadcast(conf.toJson()), batchSize));
      }