Skip to content

Commit c0c32a0

Browse files
[BugFix] LazyTensorStorage only allocates data on the given device (#2188)
1 parent d59b810 commit c0c32a0

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

torchrl/data/replay_buffers/storages.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -867,10 +867,10 @@ def max_size_along_dim0(data_shape):
867867
return (self.max_size, *data_shape)
868868

869869
if is_tensor_collection(data):
870-
out = data.expand(max_size_along_dim0(data.shape))
870+
out = data.to(self.device)
871+
out = out.expand(max_size_along_dim0(data.shape))
871872
out = out.clone()
872873
out = out.zero_()
873-
out = out.to(self.device)
874874
else:
875875
# if Tensor, we just create a MemoryMappedTensor of the desired shape, device and dtype
876876
out = tree_map(

0 commit comments

Comments
 (0)