From c35701dc713ef2c1ede8cc6d4256e0129a5c341d Mon Sep 17 00:00:00 2001
From: AdilZouitine <adilzouitinegm@gmail.com>
Date: Tue, 25 Feb 2025 15:30:39 +0000
Subject: [PATCH] Add storage device parameter to replay buffer initialization

- Specify storage device for replay buffer to optimize memory management
---
 lerobot/scripts/server/learner_server.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/lerobot/scripts/server/learner_server.py b/lerobot/scripts/server/learner_server.py
index b5c73a80..a4e42305 100644
--- a/lerobot/scripts/server/learner_server.py
+++ b/lerobot/scripts/server/learner_server.py
@@ -647,6 +647,7 @@ def train(cfg: DictConfig, out_dir: str | None = None, job_name: str | None = No
             state_keys=cfg.policy.input_shapes.keys(),
             action_mask=active_action_dims,
             action_delta=cfg.env.wrapper.delta_action,
+            storage_device=device,
         )
         batch_size: int = batch_size // 2  # We will sample from both replay buffer