Fix failing tests by catching the proper exception from torch.testing (#4121)

test/test_models.py

@@ -108,7 +108,7 @@ def _check_jit_scriptable(nn_module, args, unwrapper=None, skip=False):
         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):