Ensuring VQVAE is on the correct device when reloaded.
Ensuring VQVAE is on the correct device when reloaded.

PiperOrigin-RevId: 352005127
Change-Id: Id14d3faff6e324e16c9fd9a8ef5a8562b97d961d

deepmind/dm-haiku,deepmind/dm-haiku