update

tok/mm_autoencoder.py

@@ -12,7 +12,7 @@ class MMAutoEncoder(nn.Module):
         encoder_path, decoder_path, 
         encoder_args={}, decoder_args={}):
         super().__init__()
-        self.ar_model = {resolution: ARModel.from_checkpoint(ar_path) for resolution, ar_path in ar_path_dict.items()}
+        self.ar_model = nn.ModuleDict({resolution: ARModel.from_checkpoint(ar_path) for resolution, ar_path in ar_path_dict.items()})
 
         self.encoder = TextAlignedTokenizer.from_checkpoint(encoder_path, load_teacher=False, **encoder_args)
         self.decoder = VQVAE.from_checkpoint(decoder_path, **decoder_args)