Update iris/src/tokenizer.py

iris/src/tokenizer.py

@@ -29,9 +29,9 @@ class Tokenizer(nn.Module):
 
     def forward(self, x: torch.Tensor, should_preprocess: bool = False, should_postprocess: bool = False) -> Tuple[torch.Tensor]:
         outputs = self.encode(x, should_preprocess)
-        decoder_input = outputs.z + (outputs.z_quantized - outputs.z).detach()
+        decoder_input = outputs["z"] + (outputs["z_quantized"] - outputs["z"]).detach()
         reconstructions = self.decode(decoder_input, should_postprocess)
-        return outputs.z, outputs.z_quantized, reconstructions
+        return outputs["z"], outputs["z_quantized"], reconstructions
 
     def encode(self, x: torch.Tensor, should_preprocess: bool = False) -> dict:
         if should_preprocess: