timeout in get

pytorch · Nov 22, 2023 · 7c63518 · 7c63518
1 parent 8ca93b6
commit 7c63518
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/tensordict/base.py b/tensordict/base.py
@@ -2954,9 +2954,10 @@ def map(
             pool (mp.Pool, optional): a multiprocess Pool instance to use
                 to execute the job. If none is provided, a pool will be created
                 within the ``map`` method.
-            seed (integer, optional): the initial seed of the pool. Each member
+            generator (torch.Generator, optional): a generator to use for seeding.
+                A base seed will be generated from it, and each worker
                 of the pool will be seeded with the provided seed incremented
-                by a unique integer from ``0`` to ``num_workers``. If no seed
+                by a unique integer from ``0`` to ``num_workers``. If no generator
                 is provided, a random integer will be used as seed.
                 To work with unseeded workers, a pool should be created separately
                 and passed to :meth:`map` directly.

diff --git a/tensordict/utils.py b/tensordict/utils.py
@@ -1725,7 +1725,9 @@ def _legacy_lazy(func):
 
 # Process initializer for map
 def _proc_init(base_seed, queue):
-    worker_id = queue.get()
+    print('init worker', os.getpid())
+    worker_id = queue.get(timeout=10)
+    print('worker id', worker_id)
     seed = base_seed + worker_id
     torch.manual_seed(seed)
     np_seed = _generate_state(base_seed, worker_id)