[BugFix] Fix env.shape regex matches (pytorch#1940)

ai4co · Feb 20, 2024 · 93885b6 · 93885b6
1 parent ca42794
commit 93885b6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/test/test_env.py b/test/test_env.py
@@ -346,7 +346,7 @@ def test_mb_env_batch_lock(self, device, seed=0):
         mb_env.step(td)
 
         with pytest.raises(
-            RuntimeError, match=re.escape("Expected a tensordict with shape==env.shape")
+            RuntimeError, match=re.escape("Expected a tensordict with shape==env.batch_size")
         ):
             mb_env.step(td_expanded)
 
@@ -1615,7 +1615,7 @@ def test_batch_locked(device):
     _ = env.step(td)
 
     with pytest.raises(
-        RuntimeError, match="Expected a tensordict with shape==env.shape, "
+        RuntimeError, match="Expected a tensordict with shape==env.batch_size, "
     ):
         env.step(td_expanded)
 

diff --git a/test/test_transforms.py b/test/test_transforms.py
@@ -8231,7 +8231,7 @@ def test_batch_locked_transformed(device):
     env.step(td)
 
     with pytest.raises(
-        RuntimeError, match="Expected a tensordict with shape==env.shape, "
+        RuntimeError, match="Expected a tensordict with shape==env.batch_size, "
     ):
         env.step(td_expanded)
 
@@ -8275,7 +8275,7 @@ def test_batch_unlocked_with_batch_size_transformed(device):
     td_expanded = td.expand(2, 2).reshape(-1).to_tensordict()
 
     with pytest.raises(
-        RuntimeError, match="Expected a tensordict with shape==env.shape, "
+        RuntimeError, match="Expected a tensordict with shape==env.batch_size, "
     ):
         env.step(td_expanded)