diff --git a/rl4co/data/__init__.py b/rl4co/data/__init__.py
index e69de29b..4fa4ba8e 100644
--- a/rl4co/data/__init__.py
+++ b/rl4co/data/__init__.py
@@ -0,0 +1,2 @@
+from rl4co.data.dataset import FastTdDataset, TensorDictDataset
+from rl4co.data.utils import save_tensordict_to_npz