I guess esp in reparameterize is causing the error, as it seems to be defined without the device flag.
Could you create it with esp = torch.randn(*mu.size(), device=device) and try it again?
If that doesn’t work, could you post the whole stack trace?