Hi,
If you are trying to compute the second order partial derivative, maybe something like:
import torch

x = torch.tensor(2., requires_grad=True)
y = torch.tensor(3., requires_grad=True)

def f(a, b):
  return a**2 * b

def df_da(a, b):
  return torch.autograd.grad(f(a, b), (a,), create_graph=True)[0]

print(torch.autograd.grad(df_da(x, y), (y,))) # d^2f/dydx = 4