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

Author: Vincent Moens

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)
 

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)