fix wandb logging in tracker, and do some cleanup

dalle2_pytorch/trackers.py

@@ -1,5 +1,6 @@
 import os
 from pathlib import Path
+from enum import Enum
 import importlib
 from itertools import zip_longest
 
@@ -38,6 +39,7 @@ def load_local_state_dict(file_path, **kwargs):
 class BaseTracker(nn.Module):
     def __init__(self, data_path = DEFAULT_DATA_PATH):
         super().__init__()
+        assert data_path is not None, "Tracker must have a data_path to save local content"
         self.data_path = Path(data_path)
         self.data_path.mkdir(parents = True, exist_ok = True)
 
@@ -104,7 +106,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.log({ image_section: wandb_images }, **kwargs)
+        self.wandb.log({ image_section: wandb_images }, **kwargs)
     
     def save_state_dict(self, state_dict, relative_path, **kwargs):
         """

setup.py

@@ -10,7 +10,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.3.5',
+  version = '0.3.4',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',