edit: as to bring more detail, here is the important code I used to obtain the model:

function [dlnet] = getUncertaintyModel()

dropout_proba = 0.05;

%droplayers are not yet added, so they are commented
layers = [
    featureInputLayer(1,"Name","input");
    fullyConnectedLayer(100,"Name","fc1")
    reluLayer("Name","relu1")
    %dropoutLayer(dropout_proba,"Name","drop1")
    fullyConnectedLayer(100,"Name","fc2")
    reluLayer("Name","relu2")
    %dropoutLayer(dropout_proba,"Name","drop2")
    fullyConnectedLayer(100,"Name","fc3")
    reluLayer("Name","relu3")
    %dropoutLayer(dropout_proba,"Name","drop3")
    fullyConnectedLayer(100,"Name","fc4")
    reluLayer("Name","relu4")
    %dropoutLayer(dropout_proba,"Name","drop4")
    fullyConnectedLayer(100,"Name","fc5")
    reluLayer("Name","relu5")
    %dropoutLayer(dropout_proba,"Name","drop5")
    fullyConnectedLayer(100,"Name","fc6")
    reluLayer("Name","relu6")
    %dropoutLayer(dropout_proba,"Name","drop6")
    fullyConnectedLayer(100,"Name","fc7")
    reluLayer("Name","relu7")
    %dropoutLayer(dropout_proba,"Name","drop7")
    fullyConnectedLayer(100,"Name","fc8")
    reluLayer("Name","relu8")
    %dropoutLayer(dropout_proba,"Name","drop8")
    fullyConnectedLayer(100,"Name","fc9")
    reluLayer("Name","relu9")
    %dropoutLayer(dropout_proba,"Name","drop9")
    fullyConnectedLayer(100,"Name","fc10")
    reluLayer("Name","relu10")
    %dropoutLayer(dropout_proba,"Name","drop10")
    fullyConnectedLayer(100,"Name","fc11")
    reluLayer("Name","relu11")
    %dropoutLayer(dropout_proba,"Name","drop11")
    fullyConnectedLayer(100,"Name","fc12")
    reluLayer("Name","relu12")
    %dropoutLayer(dropout_proba,"Name","drop12")
    fullyConnectedLayer(100,"Name","fc13")
    reluLayer("Name","relu13")
    %dropoutLayer(dropout_proba,"Name","drop13")
    fullyConnectedLayer(100,"Name","fc14")
    reluLayer("Name","relu14")
    %dropoutLayer(dropout_proba,"Name","drop14")
    fullyConnectedLayer(100,"Name","fc15")
    reluLayer("Name","relu15")
    fullyConnectedLayer(1,"Name","fc16")
    
    ];

    reslgraph = layerGraph(layers);
    
    dlnet = dlnetwork(reslgraph);


end

As for the training, I use a custom training from the Matlab tutorial, and start it with this script

close all;
clear all;


%data
x= normrnd(6,1,[1,50]);
y = normrnd(cos(3.*x) ./ (abs(x) + 1.),0.03) +5;

%target
n = linspace(4,8);
n_r = cos(3.*n) ./ (abs(n) + 1.) +5;
expected = [n;n_r];

%plot
figure('Name','data');
hAxe = gca;
hold on
scatter(hAxe,x,y)
plot(hAxe,n,n_r)
hold off    

data = [x;y];
dlnet = getUncertaintyModel();

minibatchsize=20;
epochs=150;
initialLearnRate=0.01;

%start of training
 net = customTrainUncertaintyModel(dlnet,data,initialLearnRate,epochs,minibatchsize, expected);

%test data
xt=linspace(3,9);

arrxt = arrayDatastore(xt');

mbq = minibatchqueue(arrxt,...
     'MiniBatchSize',1,...
     'MiniBatchFcn',@preprocessUncertaintyMiniBatchPredictors,...
     'MiniBatchFormat',{'CB'},...
     'OutputEnvironment','auto');
 
 [YPred, means,var] = uncertaintyModelPredictions(net,mbq)
 
%plot of target/data/test prediction
figure('Name','res');
hold on
hAxe = gca;
 scatter(hAxe,x,y)
 plot(hAxe,n,n_r)
 errorbar(hAxe,xt,extractdata(means),extractdata(var),'Color',[1;0;0]);
% errorbar(hAxe,xt,means,var,'Color',[1;0;0]);
hold off

main training loop in uncertaintyModelPredictions function:

%parameters for sgdm update
velocity = [];
momentum = 0.9;

arr = arrayDatastore(data');
mbq = minibatchqueue(arr,...
     'MiniBatchSize',miniBatchSize,...
     'MiniBatchFcn',@preprocessUncertaintyMiniBatchPredictors,...
     'MiniBatchFormat',{'SSB'},...
     'OutputEnvironment','auto');
    for epoch = 1:numberOfEpochs
        shuffle(mbq);

 dsInputs = arrayDatastore(expected');
 mbqInputs = minibatchqueue(dsInputs,...
     'MiniBatchSize',miniBatchSize,...
     'MiniBatchFcn',@preprocessUncertaintyMiniBatchPredictors,...
     'MiniBatchFormat',{'SSB'},...
     'OutputEnvironment','auto');


% Loop over epochs.
for epoch = 1:numberOfEpochs
    shuffle(mbq);
    batchIndex = 0;
    
    % Loop over mini-batches.
    while hasdata(mbq)
        
        iteration = iteration + 1;
        
        batchIndex = batchIndex + 1;
        
        %extract batches from mbq
        dlData = next(mbq);
        dltempx = dlData(:,1,:);
        dltempy = dlData(:,2,:);
        dlx = dlarray(dltempx(:)','CB');
        dly = dlarray(dltempy(:)','CB');
        
       %compute gradients
       [gradients,state,loss,Ypred] = dlfeval(@uncertaintyModelGradients,dlnet,dlx,dly);
        %update state of network
        dlnet.State = state;

        learnRate = initialLearnRate;


        % Update the network parameters using the SGDM optimizer.
         [dlnet,velocity] = sgdmupdate(dlnet,gradients,velocity,learnRate,momentum);
        
        
    end
    
    reset(mbqInputs);
 
    %prediction to visualise how the model is doing on training data
    [preds] = uncertaintyModelQuickPredictions(dlnet,mbqInputs);
    scatter(haxes(2),data(1,:),data(2,:));
    hold on
    plot(haxes(2),expected(1,:),preds);
    hold off
    
    
end

and lastly here is my function to compute gradients:

function [gradients,state,loss,dlYPred] = uncertaintyModelGradients(dlnet,dlX,Y)
%compute gradients
[dlYPred,state] = forward(dlnet,dlX);
%loss
loss = mse(dlYPred,Y);

%computing
gradients = dlgradient(dlarray(loss),dlnet.Learnables);

 loss = double(gather(extractdata(loss)));

end