amend

pytorch · Jan 8, 2024 · da7933a · da7933a
1 parent ea892cc
commit da7933a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/torchrl/data/datasets/openml.py b/torchrl/data/datasets/openml.py
@@ -7,8 +7,6 @@
 
 import numpy as np
 from tensordict.tensordict import TensorDict
-
-from torchrl.data import ImmutableDatasetWriter
 from torchrl.data.replay_buffers import (
     LazyMemmapStorage,
     Sampler,
@@ -17,6 +15,8 @@
     Writer,
 )
 
+from torchrl.data.replay_buffers.writers import ImmutableDatasetWriter
+
 
 class OpenMLExperienceReplay(TensorDictReplayBuffer):
     """An experience replay for OpenML data.

diff --git a/torchrl/data/datasets/roboset.py b/torchrl/data/datasets/roboset.py
@@ -21,7 +21,7 @@
 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 Writer, ImmutableDatasetWriter
+from torchrl.data.replay_buffers.writers import ImmutableDatasetWriter, Writer
 
 _has_tqdm = importlib.util.find_spec("tqdm", None) is not None
 _has_h5py = importlib.util.find_spec("h5py", None) is not None