From 36714a14a765ef174a8e2289ac1de93ca6a603e4 Mon Sep 17 00:00:00 2001 From: AdilZouitine Date: Fri, 21 Mar 2025 14:21:31 +0000 Subject: [PATCH] Update tensor device assignment in ReplayBuffer class - Changed the device assignment for tensors in the ReplayBuffer class from `device` to `storage_device` for consistency and improved resource management. --- lerobot/scripts/server/buffer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lerobot/scripts/server/buffer.py b/lerobot/scripts/server/buffer.py index 80834eac..8ca14a03 100644 --- a/lerobot/scripts/server/buffer.py +++ b/lerobot/scripts/server/buffer.py @@ -463,9 +463,9 @@ class ReplayBuffer: for k, v in data.items(): if isinstance(v, dict): for key, tensor in v.items(): - v[key] = tensor.to(device) + v[key] = tensor.to(storage_device) elif isinstance(v, torch.Tensor): - data[k] = v.to(device) + data[k] = v.to(storage_device) action = data["action"] if action_mask is not None: