diff --git a/download_and_upload_dataset.py b/download_and_upload_dataset.py
index 2e5c806c..2d339221 100644
--- a/download_and_upload_dataset.py
+++ b/download_and_upload_dataset.py
@@ -61,7 +61,7 @@ def download_and_upload_pusht(root, root_tests, dataset_id="pusht", fps=10):
         import pymunk
         from gym_pusht.envs.pusht import PushTEnv, pymunk_to_shapely
 
-        from lerobot.common.policies.diffusion.replay_buffer import (
+        from lerobot.common.datasets._diffusion_policy_replay_buffer import (
             ReplayBuffer as DiffusionPolicyReplayBuffer,
         )
     except ModuleNotFoundError as e:
diff --git a/lerobot/common/datasets/_diffusion_policy_replay_buffer.py b/lerobot/common/datasets/_diffusion_policy_replay_buffer.py
index 1697f9fc..2f532650 100644
--- a/lerobot/common/datasets/_diffusion_policy_replay_buffer.py
+++ b/lerobot/common/datasets/_diffusion_policy_replay_buffer.py
@@ -1,6 +1,6 @@
 """Helper code for loading PushT dataset from Diffusion Policy (https://diffusion-policy.cs.columbia.edu/)
 
-Copied from the original Diffusion Policy repository.
+Copied from the original Diffusion Policy repository and used in our `download_and_upload_dataset.py` script.
 """
 
 from __future__ import annotations