diff --git a/test/test_serialization.py b/test/test_serialization.py index dcf67fe3ccf1..d12e7980c49e 100644 --- a/test/test_serialization.py +++ b/test/test_serialization.py @@ -4501,9 +4501,12 @@ class TestSerialization(TestCase, SerializationMixin): # Test that without materialize_fake_tensor, behavior for fake_tensors is not altered by ctx if not materialize_fake: ft = converter.from_real_tensor(mode, torch.randn(2, device=t_device)) + exn_type = ( + AttributeError if sys.version_info < (3, 14) else pickle.PicklingError + ) with self.assertRaisesRegex( - AttributeError, - "Can't (get|pickle) local object 'WeakValueDictionary.__init__..remove'" + exn_type, + r"Can't (get|pickle) local object .*WeakValueDictionary\.__init__\.\.remove", ): with skip_data(), BytesIOContext() as f: torch.save(ft, f)