Fixing indentation and moving test_weight_grad