Update iris/src/world_model.py

iris/src/world_model.py

@@ -27,7 +27,7 @@ class WorldModel(nn.Module):
         self.embedder = Embedder(
             max_blocks=config["max_blocks"],
             block_masks=[act_tokens_pattern, obs_tokens_pattern],
-            embedding_tables=nn.ModuleList([nn.Embedding(act_vocab_size, config["embed_dim"]), nn.Embedding(obs_vocab_size, config["embed_dim"])])
+            embedding_tables=nn.ModuleList([nn.Embedding(self.act_vocab_size, config["embed_dim"]), nn.Embedding(self.obs_vocab_size, config["embed_dim"])])
         )
 
         self.head_observations = Head(
@@ -36,7 +36,7 @@ class WorldModel(nn.Module):
             head_module=nn.Sequential(
                 nn.Linear(config["embed_dim"], config["embed_dim"]),
                 nn.ReLU(),
-                nn.Linear(config["embed_dim"], obs_vocab_size)
+                nn.Linear(config["embed_dim"], self.obs_vocab_size)
             )
         )