I have trouble understanding your intentions in this code. In particular, h_out is gradient, but it is handled the same way as regular output y_out.
In general, I’d suggest comparing empirical gradients that your custom function writes out with expected / externally precalculated values.
I’ve used this technique a bit differently; for elementwise functions Jacobian is diagonal and extra grad() if backward is unnecessary. It looked like:
#in forward
dY_dX = torch.autograd.grad(y, x_autograd, y.new_ones(1).expand_as(y))
ctx.save_for_backward(*dY_dX)
#in backward
known_grads = ctx.saved_tensors
grad_inputs[idx] = grad_output * known_grads[idx]
return (*grad_inputs,)

Hence, dY_dX tensors are plugged into backward(). For other functions, this approach would require autograd.jacobian() (likely prohibitive) and matrix multiplications.