How to calculate diagonal elements of a hessian matrix in pytorch

Hi,
I am trying to calculate diagonal elements of a hessian matrix. I used the below code:

g—>> 1D tensor of gradients.
diag=[]
for i in range(len(g)):
cnn.zero_grad()
g[i].backward(retain_graph=True, create_graph=False, inputs=p)
ele=p[i].grad
diag.append(ele)

But it is taking long time.

Any suggestions how to calculate diagonal elements of a hessian in an efficient way.
Thanks in advance