diff --git a/pokemonred_puffer/cleanrl_puffer.py b/pokemonred_puffer/cleanrl_puffer.py
index 95827c9..ec21099 100644
--- a/pokemonred_puffer/cleanrl_puffer.py
+++ b/pokemonred_puffer/cleanrl_puffer.py
@@ -302,7 +302,7 @@ def evaluate(self):
 
                 for k, v in self.infos.items():
                     if "_map" in k and self.wandb_client is not None:
-                        self.stats[f"Media/{k}"] = self.wandb.Image(v[0])
+                        self.stats[f"Media/{k}"] = self.wandb_client.Image(v[0])
                         continue
                     elif "state" in k:
                         pass
diff --git a/pokemonred_puffer/train.py b/pokemonred_puffer/train.py
index 164d006..a37f6e8 100644
--- a/pokemonred_puffer/train.py
+++ b/pokemonred_puffer/train.py
@@ -289,7 +289,7 @@ def train(
 
     wandb_client = None
     if args.track:
-        wandb_client = init_wandb(args).id
+        wandb_client = init_wandb(args)
 
     if args.mode == "train":
         train(args, env_creator, wandb_client)