diff --git a/pokemonred_puffer/cleanrl_puffer.py b/pokemonred_puffer/cleanrl_puffer.py
index c8230a3..73c5e2f 100644
--- a/pokemonred_puffer/cleanrl_puffer.py
+++ b/pokemonred_puffer/cleanrl_puffer.py
@@ -303,11 +303,9 @@ def evaluate(self):
             for k, v in self.infos.items():
                 # Moves into models... maybe. Definitely moves.
                 # You could also just return infos and have it in demo
-                if "pokemon_exploration_map" in self.infos and self.config.save_overlay is True:
+                if "pokemon_exploration_map" in k and self.config.save_overlay is True:
                     if self.epoch % self.config.overlay_interval == 0:
-                        overlay = make_pokemon_red_overlay(
-                            np.stack(self.infos["pokemon_exploration_map"], axis=0)
-                        )
+                        overlay = make_pokemon_red_overlay(np.stack(self.infos[k], axis=0))
                         if self.wandb_client is not None:
                             self.stats["Media/aggregate_exploration_map"] = wandb.Image(
                                 overlay, file_type="jpg"