function data = CollateNDData(cellData, outputLayerSize, N)
% Convert BatchSize x NumVariables batches into 1xNumVariables cell
% array where each column contains a N-D array or Tensor object.

%   Copyright 2018-2021 The MathWorks, Inc.

data = cell([1,size(cellData,2)]);
responseCol = 2;

for col = 1:size(cellData,2)
    if iscategorical(cellData{1,col})
        % Detect if the observation is scalar
        isScalarObs = isscalar(cellData{1,1});
        if isvector(cellData{1,col})
            % Vector response for classification.
            responses = [cellData{:,col}];
        else
            % Tensor response for semantic segmentation.
            responses = cat(N,cellData{:,col});
        end
        data{1,col} = nnet.internal.cnn.util.dummifyND(responses, N, isScalarObs);
    else
        
        % The response needs to be reshaped for regression only for scalar
        % or vector regression with outputLayerSize = [1 1 K] or [1 1 1 K]. 
        % No reshape is needed for image to image regression case.
        reshapeNeeded = (prod(outputLayerSize(1:end-1)) == 1) &&...
                                    (col == responseCol);
        % Numeric data
        if reshapeNeeded
            % Reshape numeric responses for regression
            reshapedCell = cellfun(@(x)reshape(x,outputLayerSize), cellData(:,col), ...
                                            'UniformOutput', false);
        else
            reshapedCell = cellData(:,col);
        end
        data{1,col} = single(cat(N, reshapedCell{:}));
    end
end

end
