Implement mean squared error over all elements and its gradient with respect to the predictions:
mse(pred, target)
mse_grad(pred, target)
pred