Missing dependency and problem with import (#272)

requirements.txt

@@ -8,6 +8,7 @@ syne-tune[aws,gpsearchers]==0.6.0
 pytorch-lightning>=1.8.0, <1.9.5
 Pillow>=9.0, <9.5.1
 tabulate>=0.9.0, <0.9.1
+tensorboardX>=2.5.0, <2.5.2
 torchmetrics~=0.10.3
 torchvision>=0.13.0, <0.15.2
 deepspeed==0.9.1

src/renate/benchmark/datasets/wild_time_data.py

@@ -4,8 +4,6 @@
 from typing import Any, Dict, Optional, Union
 
 from transformers import PreTrainedTokenizer
-from wild_time_data import load_dataset
-from wild_time_data.core import available_time_steps, dataset_classes
 
 from renate import defaults
 from renate.data.data_module import RenateDataModule
@@ -63,13 +61,17 @@ def prepare_data(self) -> None:
         If s3 bucket is given, the data is downloaded from s3, otherwise from the original source.
         """
         if self._src_bucket is None:
+            from wild_time_data import available_time_steps, load_dataset
+
             load_dataset(
                 dataset_name=self._dataset_name,
                 time_step=available_time_steps(self._dataset_name)[0],
                 split="train",
                 data_dir=self._data_path,
             )
         else:
+            from wild_time_data.core import dataset_classes
+
             dst_dir = Path(self._data_path) / dataset_classes[self._dataset_name].file_name
             if not dst_dir.exists():
                 download_folder_from_s3(
@@ -80,6 +82,8 @@ def prepare_data(self) -> None:
 
     def setup(self) -> None:
         """Set up train, test and val datasets."""
+        from wild_time_data import available_time_steps, load_dataset
+
         kwargs = {
             "dataset_name": self._dataset_name,
             "time_step": available_time_steps(self._dataset_name)[self.time_step],