Skip to content

Commit 313812d

Browse files
committed
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.
1 parent 3692f88 commit 313812d

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

Diff for: lerobot/scripts/server/buffer.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -463,9 +463,9 @@ def from_lerobot_dataset(
463463
for k, v in data.items():
464464
if isinstance(v, dict):
465465
for key, tensor in v.items():
466-
v[key] = tensor.to(device)
466+
v[key] = tensor.to(storage_device)
467467
elif isinstance(v, torch.Tensor):
468-
data[k] = v.to(device)
468+
data[k] = v.to(storage_device)
469469

470470
action = data["action"]
471471
if action_mask is not None:

0 commit comments

Comments
 (0)