Hi,
The thing is that d(y_pred)/d(z) is a 2D Jacobian matrix. The autograd engine can only compute Jacobian vector product.
Could you give more informations about what func() is? Maybe it is doable in your case.