[BugFix] Fix env.shape regex matches

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)