I've made a working example on colab: https://colab.research.google.com/gist/philastrophist/e7caf81960b6a6898b68e7a0f83b524c/svi_compare.ipynb
Basically, there is some sort of confusion between cuda/cpu data loaders. I don't know how to fix this myself.
Thanks