You are returning z which is your input and as @mruberry has mentioned, if you need grad for your inputs, then you should enable it. But you have mentioned that you only need grads w.r.t. model parameters (not input). In that case, you should return samples instead of z to see grads. I think you are returning the wrong tensor.