amend

pytorch · Feb 7, 2025 · 208f087 · 208f087
1 parent 7ce5263
commit 208f087
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tutorials/sphinx-tutorials/pretrained_models.py b/tutorials/sphinx-tutorials/pretrained_models.py
@@ -37,7 +37,7 @@
 import torch.cuda
 from tensordict.nn import TensorDictSequential
 from torch import nn
-from torchrl.envs import R3MTransform, TransformedEnv
+from torchrl.envs import R3MTransform, TransformedEnv, Compose
 from torchrl.envs.libs.gym import GymEnv
 from torchrl.modules import Actor
 
@@ -115,7 +115,7 @@
 from torchrl.data import LazyMemmapStorage, ReplayBuffer
 
 storage = LazyMemmapStorage(1000)
-rb = ReplayBuffer(storage=storage, transform=r3m)
+rb = ReplayBuffer(storage=storage, transform=Compose(lambda td: td.to(device), r3m))
 
 ##############################################################################
 # We can now collect the data (random rollouts for our purpose) and fill the replay