Update iris/src/world_model.py

iris/src/world_model.py

@@ -12,7 +12,7 @@ from .models.transformer import Transformer
 class WorldModel(nn.Module):
     def __init__(self, config: dict) -> None:
         super().__init__()
-        self.obs_vocab_size, self.act_vocab_size = config["obs_vocab_size"], config["act_vocab_size"]
+        self.obs_vocab_size, self.act_vocab_size = config["vocab_size"], config["act_vocab_size"]
         self.config = config
         self.transformer = Transformer(config)