There are two types of data: numpy and torch

z = mean.clone().detach() + eps * torch.exp(logstd)

An error is reported in the source code, which is modified as follows

eps = eps.cuda()    
z = mean.cuda()+ eps * torch.exp(logstd).cuda()

It is to add CUDA() after numpy, such as reshape