Could you use the regression loss functions from the SGD module (squared, huber and epsilon-insensitive). This is more flexible and this way we are sure that they updates are correct.