Update iris/src/world_model.py

iris/src/world_model.py

@@ -84,9 +84,9 @@ class WorldModel(nn.Module):
 
         }
 
-    def generate_empty_keys_values(self):
+    def generate_empty_keys_values(self, n= 1):
 
-        values = self.world_model.transformer.generate_empty_keys_values(n=1, max_tokens= self.config["max_tokens"])
+        values = self.world_model.transformer.generate_empty_keys_values(n=n, max_tokens= self.config["max_tokens"])
     
     def compute_labels_world_model(self, obs_tokens: torch.Tensor, rewards: torch.Tensor, ends: torch.Tensor, mask_padding: torch.BoolTensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         assert torch.all(ends.sum(dim=1) <= 1)  # at most 1 done