Update

[ghstack-poisoned]
pytorch · Feb 20, 2025 · 5e8c76f · 5e8c76f
1 parent 43b735b
commit 5e8c76f
Showing 1 changed file with 13 additions and 12 deletions.
diff --git a/test/test_env.py b/test/test_env.py
@@ -2996,18 +2996,19 @@ def test_parallel(
     ],
 )
 def test_mocking_envs(envclass):
-    env = envclass()
-    with pytest.warns(UserWarning, match="model based") if isinstance(
-        env, DummyModelBasedEnvBase
-    ) else contextlib.nullcontext():
-        env.set_seed(100)
-    reset = env.reset()
-    _ = env.rand_step(reset)
-    r = env.rollout(3)
-    with pytest.warns(UserWarning, match="model based") if isinstance(
-        env, DummyModelBasedEnvBase
-    ) else contextlib.nullcontext():
-        check_env_specs(env, seed=100, return_contiguous=False)
+    with set_capture_non_tensor_stack(False):
+        env = envclass()
+        with pytest.warns(UserWarning, match="model based") if isinstance(
+            env, DummyModelBasedEnvBase
+        ) else contextlib.nullcontext():
+            env.set_seed(100)
+        reset = env.reset()
+        _ = env.rand_step(reset)
+        r = env.rollout(3)
+        with pytest.warns(UserWarning, match="model based") if isinstance(
+            env, DummyModelBasedEnvBase
+        ) else contextlib.nullcontext():
+            check_env_specs(env, seed=100, return_contiguous=False)
 
 
 class TestTerminatedOrTruncated: