[dynamo] wrap GraphModule exceptions in dynamo-wrapped tests (#126341)

Better approach to https://github.com/pytorch/pytorch/pull/126197 to catch issues like https://github.com/pytorch/pytorch/issues/125568.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/126341
Approved by: https://github.com/anijain2305, https://github.com/jansel
This commit is contained in:
William Wen
2024-05-28 14:05:39 -07:00
committed by PyTorch MergeBot
parent cdf2133186
commit 5359af0c7e
18 changed files with 101 additions and 9 deletions

View File

@ -77,6 +77,7 @@ from torch.testing._internal.common_utils import (
subtest,
TEST_WITH_TORCHDYNAMO,
TestCase,
xfailIfTorchDynamo,
)
from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
@ -2340,6 +2341,8 @@ class TestJac(VmapTearDownMixin, TestCase):
)(x)
self.assertEqual(actual, expected)
# https://github.com/pytorch/pytorch/issues/127036
@xfailIfTorchDynamo
@parametrize("_preallocate_and_copy", (True, False))
def test_chunk_jacrev_chunksize_one(self, device, _preallocate_and_copy):
# With chunk_size=1, we shouldn't `vmap` and hence not be limited