How to train CEVAE with GPU?

[xiaogangzhu]

Is your feature request related to a problem? Please describe.
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
I find that CEVAE is very useful for me, but unfortunately I found it is not training using GPU and I cannot modify the dataloader to imporove my data reading efficiency.

Describe the solution you'd like
A clear and concise description of what you want to happen.
Is there any way to train CEVAE with GPU and also improve the Dataloader to have multiple workers?

Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.

Additional context
Add any other context or screenshots about the feature request here.

[xiaogangzhu]

Any update?