amend

pytorch · Nov 24, 2023 · ed5b511 · ed5b511
1 parent c5b4bda
commit ed5b511
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/tutorials/sphinx-tutorials/rb_tutorial.py b/tutorials/sphinx-tutorials/rb_tutorial.py
@@ -46,6 +46,7 @@
 # replay buffer is a straightforward process, as shown in the following
 # example:
 #
+import tempfile
 
 from torchrl.data import ReplayBuffer
 
@@ -175,8 +176,9 @@
 ######################################################################
 # We can also customize the storage location on disk:
 #
+tempdir = tempfile.TemporaryDirectory()
 buffer_lazymemmap = ReplayBuffer(
-    storage=LazyMemmapStorage(size, scratch_dir="/tmp/memmap/")
+    storage=LazyMemmapStorage(size, scratch_dir=tempdir)
 )
 buffer_lazymemmap.extend(data)
 print(f"The buffer has {len(buffer_lazymemmap)} elements")
@@ -207,8 +209,9 @@
 
 from torchrl.data import TensorDictReplayBuffer
 
+tempdir = tempfile.TemporaryDirectory()
 buffer_lazymemmap = TensorDictReplayBuffer(
-    storage=LazyMemmapStorage(size, scratch_dir="/tmp/memmap/"), batch_size=12
+    storage=LazyMemmapStorage(size, scratch_dir=tempdir), batch_size=12
 )
 buffer_lazymemmap.extend(data)
 print(f"The buffer has {len(buffer_lazymemmap)} elements")
@@ -248,8 +251,9 @@ class MyData:
     batch_size=[1000],
 )
 
+tempdir = tempfile.TemporaryDirectory()
 buffer_lazymemmap = TensorDictReplayBuffer(
-    storage=LazyMemmapStorage(size, scratch_dir="/tmp/memmap/"), batch_size=12
+    storage=LazyMemmapStorage(size, scratch_dir=tempdir), batch_size=12
 )
 buffer_lazymemmap.extend(data)
 print(f"The buffer has {len(buffer_lazymemmap)} elements")