It didn’t work even after you detach the variables? How about directly using the underlying tensors, i.e., z.data and zb.data to index?
FYI: what you are observing is a bug recently exposed https://github.com/pytorch/pytorch/issues/3348 6