fix vae nan bug

vqvae.py

@@ -58,6 +58,7 @@ class Autoencoder(nn.Module):
         # Get embeddings with shape `[batch_size, z_channels * 2, z_height, z_height]`
         print(f"encoder parameters max: {max([p.max() for p in self.encoder.parameters()])}")
         print(f"encoder parameters min: {min([p.min() for p in self.encoder.parameters()])}")
+        print(f"img.dtype: {img.dtype}")
         z = self.encoder(img)
         print(f"z.max(): {z.max()}, z.min(): {z.min()}")
         # Get the moments in the quantized embedding space