[learning] Fix loss

Current loss only compute loss but it should also compute the output gradient of the last layer.

To change this I will make so that loss function take a graph and the ground truth labels as argument and return the loss.