From 7caddc3cf988b747c1e6f8f69295cea2640ae4e8 Mon Sep 17 00:00:00 2001
From: vmoens <vincentmoens@gmail.com>
Date: Tue, 9 Jan 2024 09:03:48 +0000
Subject: [PATCH] amend

---
 torchrl/data/datasets/d4rl.py   | 2 +-
 torchrl/data/datasets/openml.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/torchrl/data/datasets/d4rl.py b/torchrl/data/datasets/d4rl.py
index eb84cf380d2..20667c33525 100644
--- a/torchrl/data/datasets/d4rl.py
+++ b/torchrl/data/datasets/d4rl.py
@@ -24,7 +24,7 @@
 from torchrl.data.datasets.d4rl_infos import D4RL_DATASETS
 
 from torchrl.data.datasets.utils import _get_root_dir
-from torchrl.data.replay_buffers import TensorDictReplayBuffer
+from torchrl.data.replay_buffers.replay_buffers import TensorDictReplayBuffer
 from torchrl.data.replay_buffers.samplers import Sampler
 from torchrl.data.replay_buffers.storages import TensorStorage
 from torchrl.data.replay_buffers.writers import ImmutableDatasetWriter, Writer
diff --git a/torchrl/data/datasets/openml.py b/torchrl/data/datasets/openml.py
index e351e94cf4e..feac3ec8b0d 100644
--- a/torchrl/data/datasets/openml.py
+++ b/torchrl/data/datasets/openml.py
@@ -9,12 +9,12 @@
 import numpy as np
 from tensordict.tensordict import TensorDict
 
-from torchrl.data import TensorStorage
 from torchrl.data.datasets.utils import _get_root_dir
 from torchrl.data.replay_buffers import (
     Sampler,
     SamplerWithoutReplacement,
     TensorDictReplayBuffer,
+    TensorStorage,
     Writer,
 )