Stats plot for loss function of cnn_train_dag

[Sidra128]

While plotting the loss value at the end of the epoch in cnn_train_dag function, we plot only the loss value of the last training batch. The function assigns the value of the loss computed for last batch to the variable stats, that is plotted at the end of each epoch

Wouldn't it better if we plot the average value of the loss computed for all the batches processed during an epoch? This would better show the progress of learning in subsequent epochs, as we do not know that the randomly sampled batches might have performed best or worst for the last batch of an epoch.