Update iris/src/world_model.py

iris/src/world_model.py

@@ -83,6 +83,10 @@ class WorldModel(nn.Module):
             "logits_ends": logits_ends
 
         }
+
+    def generate_empty_keys_values(self):
+
+        values = self.world_model.transformer.generate_empty_keys_values(n=1, 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