Optimizing an autograd + list comprehension python operation

Question

I have two tensors M (N, N) and D (N, T, F), and I want to calculate the derivative of M[i, j] with respect to D[j, -1]. At the moment I'm using torch.autograd.grad(M[i, j], D) and getting the [j, -1]th output, but I'm wondering if it's possible to differentiate directly with respect to D[j, -1].

The final goal is to calculate this:

torch.stack([
    torch.stack([
        torch.autograd.grad(K[i, j], X, retain_graph=True)[0][j, -1]
        for j in range(X.shape[0])
    ])
    for i in range(X.shape[0])
]).mean(1)

If you can think of a way to optimise it, I'd love to hear from you!