Update iris/src/tokenizer.py

iris/src/tokenizer.py

@@ -16,11 +16,11 @@ class Tokenizer(nn.Module):
         super().__init__()
         self.vocab_size = config["vocab_size"]
         self.embed_dim = config["embed_dim"]
-        self.encoder = Encoder(config)
+        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)
+        self.decoder = Decoder(config["decoder"])
         self.embedding.weight.data.uniform_(-1.0 / self.vocab_size, 1.0 / self.vocab_size)
         self.lpips = LPIPS().eval() if with_lpips else None