From e74f41efc6ddea1df424997335cfc1a64dac8bd3 Mon Sep 17 00:00:00 2001
From: Nicolas Hug <nicolashug@fb.com>
Date: Fri, 25 Jun 2021 14:54:18 +0100
Subject: [PATCH] change error

---
 test/test_models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/test_models.py b/test/test_models.py
index d40b72e55ff..eb8245f3ad4 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 pytest.UsageError:
+        except torch.testing._asserts.UsageError:
             # 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):