Refactor loss function and gradient compilation.