1. Create iterators for the MNIST data using MnistDataSetIterator: 

DataSetIterator iter = new MnistDataSetIterator(miniBatchSize,numOfExamples,binarize);

2. Use SplitTestAndTrain to split the base iterator into train/test iterators:

DataSet ds = iter.next();
 SplitTestAndTrain split = ds.splitTestAndTrain(numHoldOut, new Random(12345));

3. Create lists to store the feature sets from the train/test iterators:

List<INDArray> featuresTrain = new ArrayList<>();
 List<INDArray> featuresTest = new ArrayList<>();
 List<INDArray> labelsTest = new ArrayList<>();

4. Populate the values into the feature/label lists that were previously created:

featuresTrain.add(split.getTrain().getFeatures());
 DataSet dsTest = split.getTest();
 featuresTest.add(dsTest.getFeatures());
 INDArray indexes = Nd4j.argMax(dsTest.getLabels(),1);
 labelsTest.add(indexes);

5. Call argmax() for every iterator instance to convert the labels to one dimensional data if it's multidimensional:

while(iter.hasNext()){
 DataSet ds = iter.next();
 SplitTestAndTrain split = ds.splitTestAndTrain(80, new Random(12345)); // 80/20 split (from miniBatch = 100)
 featuresTrain.add(split.getTrain().getFeatures());
 DataSet dsTest = split.getTest();
 featuresTest.add(dsTest.getFeatures());
 INDArray indexes = Nd4j.argMax(dsTest.getLabels(),1);
 labelsTest.add(indexes);
 }