Update iris/src/tokenizer.py

iris/src/tokenizer.py

@@ -18,9 +18,9 @@ class Tokenizer(nn.Module):
         self.embed_dim = config["embed_dim"]
         self.encoder = Encoder(config["encoder"])
         self.decoder = Decoder(config["decoder"])
-        self.pre_quant_conv = torch.nn.Conv2d(self.encoder.config.z_channels, self.embed_dim, 1)
+        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.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