amend

pytorch · Nov 23, 2023 · 401f055 · 401f055
1 parent 2fa557f
commit 401f055
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/tensordict/utils.py b/tensordict/utils.py
@@ -1726,9 +1726,7 @@ def _legacy_lazy(func):
 
 # Process initializer for map
 def _proc_init(base_seed, queue):
-    print("worker", os.getpid())
     worker_id = queue.get(timeout=10)
-    print(worker_id)
     seed = base_seed + worker_id
     torch.manual_seed(seed)
     np_seed = _generate_state(base_seed, worker_id)

diff --git a/test/test_tensordict.py b/test/test_tensordict.py
@@ -6557,9 +6557,9 @@ class TestMap:
     @classmethod
     def get_rand_incr(cls, td):
         # torch
-        td["r"] += torch.randint(0, 100, ()).item()
+        td["r"] = td["r"] + torch.randint(0, 100, ()).item()
         # numpy
-        td["s"] += np.random.randint(0, 100, ()).item()
+        td["s"] = td["s"] + np.random.randint(0, 100, ()).item()
         return td
 
     def test_map_seed(self):
@@ -6572,7 +6572,7 @@ def test_map_seed(self):
                 "c": torch.arange(20),
             },
             batch_size=[20],
-        ).memmap_()
+        )
         generator = torch.Generator()
         # we use 4 workers with max 5 items each,
         # making sure that no worker does more than any other.