fix vae nan bug

vqvae.py

@@ -56,9 +56,13 @@ class Autoencoder(nn.Module):
         :param img: is the image tensor with shape `[batch_size, img_channels, img_height, img_width]`
         """
         # 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()])}")
         z = self.encoder(img)
+        print(f"z.max(): {z.max()}, z.min(): {z.min()}")
         # Get the moments in the quantized embedding space
         moments = self.quant_conv(z)
+        print(f"moments.max(): {moments.max()}, moments.min(): {moments.min()}")
         # Return the distribution
         return GaussianDistribution(moments)
 
@@ -284,6 +288,7 @@ class GaussianDistribution:
             `[batch_size, z_channels * 2, z_height, z_height]`
         """
         # Split mean and log of variance
+        print(f"parameters.max(): {parameters.max()}, parameters.min(): {parameters.min()}")
         self.mean, log_var = torch.chunk(parameters, 2, dim=1)
         # Clamp the log of variances
         self.log_var = torch.clamp(log_var, -30.0, 20.0)
@@ -293,6 +298,8 @@ class GaussianDistribution:
 
     def sample(self):
         # Sample from the distribution
+        print(f"self.mean.max(): {self.mean.max()}, self.mean.min(): {self.mean.min()}")
+        print(f"self.std.max(): {self.std.max()}, self.std.min(): {self.std.min()}")
         return self.mean + self.std * torch.randn_like(self.std)
 
     def kl(self):