Add gradient checking.
Minor changes