diff --git a/README.rst b/README.rst
index a707773c..30f6028e 100644
--- a/README.rst
+++ b/README.rst
@@ -2,7 +2,7 @@
     :target: #
     :alt: PyPI - Status
 .. image:: https://img.shields.io/github/v/release/awslabs/Renate
-    :target: https://github.com/awslabs/Renate/releases/tag/v0.2.0
+    :target: https://github.com/awslabs/Renate/releases/tag/v0.3.1
     :alt: Latest Release
 .. image:: https://img.shields.io/pypi/dm/Renate
     :target: https://pypistats.org/packages/renate
diff --git a/requirements.txt b/requirements.txt
index c9ce9bc9..11ca7b19 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -8,8 +8,9 @@ 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
-datasets>=2.9.0, < 2.12.1
+datasets>=2.9.0, <2.12.1
 transformers>4.23.0, <4.29.3
 scipy>=1.9.0, <1.10.2
diff --git a/src/renate/__init__.py b/src/renate/__init__.py
index 5deff2eb..a6ccdc9d 100644
--- a/src/renate/__init__.py
+++ b/src/renate/__init__.py
@@ -14,4 +14,4 @@
     _renate_logger.addHandler(_handler)
     _renate_logger.propagate = False
 
-__version__ = "0.3.0"
+__version__ = "0.3.1"
diff --git a/src/renate/benchmark/datasets/wild_time_data.py b/src/renate/benchmark/datasets/wild_time_data.py
index d1727c5f..2655ae85 100644
--- a/src/renate/benchmark/datasets/wild_time_data.py
+++ b/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,6 +61,8 @@ 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],
@@ -70,6 +70,8 @@ def prepare_data(self) -> None:
                 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],