I have checked the inputs and found some NaNs so I found the responsible layer which was a nn.ConvTranspose2d(embed_dim, embed_dim, 4, stride=2, padding=1 ) so I made it run in float32 and now there’s no more NaNs but the loss explodes to values like 6527757.50000 . So again, I find the responsible layer, self.quantize_t = Quantize(embed_dim, n_embed) but when I try to make it run in float32 I get the error:
\vqvae.py", line 46, in forward
+ self.embed.pow(2).sum(0, keepdim=True)
RuntimeError: expected scalar type Half but found Float
during the forward pass of the layer.
At this point, im not sure this is even feasible.