This is slightly slower on CPU…
grads = torch.autograd.grad(loss, weights, create_graph=True)[0] # size (D)
grads_times_Z = torch.mm(grads.unsqueeze(1).T, Z)
for k in range(K):
    HZp[:,k] = torch.autograd.grad(grads_times_Z[:,k], weights, retain_graph=True)[0]