Update delta-iris/src/world_model.py

delta-iris/src/world_model.py

@@ -58,7 +58,7 @@ class WorldModel(nn.Module):
         return self.transformer.num_blocks_left_in_kv_cache
 
     def reset_kv_cache(self):
-        transformer.reset_kv_cache(n=1)
+        self.transformer.reset_kv_cache(n=1)
 
     def forward(self, sequence: torch.FloatTensor, use_kv_cache: bool = False) -> dict:      
         prev_steps = self.transformer.keys_values.size if use_kv_cache else 0