cleaned up some of the gradient checking