test_single_trans_env_check

pytorch · Jan 4, 2024 · a7b72e9 · a7b72e9
1 parent df3edba
commit a7b72e9
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/test/test_transforms.py b/test/test_transforms.py
@@ -9407,14 +9407,15 @@ def _make_batch(self, batch_size: int = 2, sequence_length: int = 5):
         return batch
 
     def test_single_trans_env_check(self):
-        # raise pytest.skip(
-        #         "BurnInTransform can only be appended to a ReplayBuffer, not to a TransformedEnv."
-        #     )
         module = self._make_gru_module()
         burn_in_transform = BurnInTransform(module, burn_in=2)
-        env = TransformedEnv(ContinuousActionVecMockEnv(), burn_in_transform)
-        check_env_specs(env)
-        env.close()
+        with pytest.raises(
+                RuntimeError,
+                match="BurnInTransform can only be appended to a ReplayBuffer.",
+        ):
+            env = TransformedEnv(ContinuousActionVecMockEnv(), burn_in_transform)
+            check_env_specs(env)
+            env.close()
 
     def test_serial_trans_env_check(self):
         raise pytest.skip(