diff --git a/dalle2_pytorch/trackers.py b/dalle2_pytorch/trackers.py
index a5dbe84..9ae56ff 100644
--- a/dalle2_pytorch/trackers.py
+++ b/dalle2_pytorch/trackers.py
@@ -1,6 +1,5 @@
 import os
 from pathlib import Path
-from enum import Enum
 import importlib
 from itertools import zip_longest
 
@@ -105,7 +104,7 @@ class WandbTracker(BaseTracker):
         Takes a tensor of images and a list of captions and logs them to wandb.
         """
         wandb_images = [self.wandb.Image(image, caption=caption) for image, caption in zip_longest(images, captions)]
-        self.wandb.log({ image_section: wandb_images }, **kwargs)
+        self.log({ image_section: wandb_images }, **kwargs)
     
     def save_state_dict(self, state_dict, relative_path, **kwargs):
         """
diff --git a/setup.py b/setup.py
index 638b78a..4a191d7 100644
--- a/setup.py
+++ b/setup.py
@@ -10,7 +10,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.3.4',
+  version = '0.3.5',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',