diff --git a/test/test_models.py b/test/test_models.py index ee4af8bef58..72ae68f5615 100644 --- a/test/test_models.py +++ b/test/test_models.py @@ -108,7 +108,7 @@ def get_export_import_copy(m): tol = 3e-4 try: torch.testing.assert_close(results, results_from_imported, atol=tol, rtol=tol) - except torch.testing._asserts.UsageError: + except ValueError: # custom check for the models that return named tuples: # we compare field by field while ignoring None as assert_close can't handle None for a, b in zip(results, results_from_imported):