Update delta-iris/src/models/tokenizer/quantizer.py

delta-iris/src/models/tokenizer/quantizer.py

@@ -20,7 +20,7 @@ class Quantizer(nn.Module):
         self.register_buffer('codebook', codebook)
         self.register_buffer('codewords_freqs', torch.ones(codebook_size).div(codebook_size))
 
-    def forward(self, z: torch.Tensor) -> QuantizerOutput:
+    def forward(self, z: torch.Tensor) -> dict:
         z = self.pre_quant_proj(z)
         z = F.normalize(z, dim=-1)
         b, k = z.size(0), z.size(2)