minor changes to check_grad output