Most likely this error is thrown because esp is still on the CPU while mu and std seem to be calculated by the linear layers and might thus be on the GPU (if the model was pushed to the GPU of course).
Try to create esp using the device argument of mu:
esp = torch.randn(*mu.size(), device=mu.device)