[Test] Capture deprec warnings

ghstack-source-id: bcbf41c245c979d0f21524889ad2be8ef4c10c40 Pull Request resolved: #2799
pytorch · Feb 20, 2025 · fb641de · fb641de
1 parent ea76ffb
commit fb641de
Show file tree

Hide file tree

Showing 9 changed files with 468 additions and 388 deletions.
diff --git a/test/_utils_internal.py b/test/_utils_internal.py
@@ -9,6 +9,7 @@
 import os
 
 import os.path
+import sys
 import time
 import unittest
 import warnings
@@ -42,12 +43,17 @@
     ToTensorImage,
     TransformedEnv,
 )
+from torchrl.modules import MLP
 from torchrl.objectives.value.advantages import _vmap_func
 
 # Specified for test_utils.py
 __version__ = "0.3"
 
-from torchrl.modules import MLP
+IS_WIN = sys.platform == "win32"
+if IS_WIN:
+    mp_ctx = "spawn"
+else:
+    mp_ctx = "fork"
 
 
 def CARTPOLE_VERSIONED():
@@ -265,6 +271,7 @@ def _make_envs(
     N,
     device="cpu",
     kwargs=None,
+    local_mp_ctx=mp_ctx,
 ):
     torch.manual_seed(0)
     if not transformed_in:
@@ -299,7 +306,9 @@ def create_env_fn():
                 )
 
     env0 = create_env_fn()
-    env_parallel = ParallelEnv(N, create_env_fn, create_env_kwargs=kwargs)
+    env_parallel = ParallelEnv(
+        N, create_env_fn, create_env_kwargs=kwargs, mp_start_method=local_mp_ctx
+    )
     env_serial = SerialEnv(N, create_env_fn, create_env_kwargs=kwargs)
 
     for key in env0.observation_spec.keys(True, True):