Fix test_batchsize_reset

pytorch · Feb 9, 2023 · 7a05d2c · 7a05d2c
1 parent 6fe9382
commit 7a05d2c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
@@ -2489,8 +2489,8 @@ def test_batchsize_reset():
     # test index
     td[torch.tensor([1, 2])]
     with pytest.raises(
-        IndexError,
-        match=re.escape("too many indices for tensor of dimension 1"),
+        RuntimeError,
+        match=re.escape("The shape torch.Size([3]) is incompatible with the index"),
     ):
         td[:, 0]