Update iris/src/tokenizer.py

iris/src/tokenizer.py

@@ -17,10 +17,10 @@ class Tokenizer(nn.Module):
         self.vocab_size = config["vocab_size"]
         self.embed_dim = config["embed_dim"]
         self.encoder = Encoder(config["encoder"])
-        self.pre_quant_conv = torch.nn.Conv2d(encoder.config.z_channels, self.embed_dim, 1)
-        self.embedding = nn.Embedding(self.vocab_size, self.embed_dim)
-        self.post_quant_conv = torch.nn.Conv2d(self.embed_dim, decoder.config.z_channels, 1)
         self.decoder = Decoder(config["decoder"])
+        self.pre_quant_conv = torch.nn.Conv2d(self.encoder.config.z_channels, self.embed_dim, 1)
+        self.embedding = nn.Embedding(self.vocab_size, self.embed_dim)
+        self.post_quant_conv = torch.nn.Conv2d(self.embed_dim, self.decoder.config.z_channels, 1)
         self.embedding.weight.data.uniform_(-1.0 / self.vocab_size, 1.0 / self.vocab_size)
         self.lpips = LPIPS().eval() if with_lpips else None