Replace assertRaisesRegex w/ assertRaisesRegexWithHighlight test_backends.py (#55493)

Summary:
Step to resolving https://github.com/pytorch/pytorch/issues/55072

Pull Request resolved: https://github.com/pytorch/pytorch/pull/55493

Reviewed By: gmagogsfm

Differential Revision: D27626192

Pulled By: janeyx99

fbshipit-source-id: 047b7b6754e21388f52489160d712858b7aa0288
diff --git a/test/jit/test_backends.py b/test/jit/test_backends.py
index 2e1b786..0f67cd0 100644
--- a/test/jit/test_backends.py
+++ b/test/jit/test_backends.py
@@ -196,7 +196,9 @@
         input = torch.randn(5)
 
         # Test exception is thrown.
-        with self.assertRaisesRegex(Exception, r"Backend is not available."):
+        with self.assertRaisesRegexWithHighlight(Exception,
+                                                 r"Backend is not available.",
+                                                 "raise Exception(\"Backend is not available.\""):
             backend_method = self.lowered_module.__getattr__("forward")
             backend_output = backend_method(*(input, input))
 
@@ -206,7 +208,9 @@
         buffer = io.BytesIO()
         torch.jit.save(self.lowered_module, buffer)
         buffer.seek(0)
-        with self.assertRaisesRegex(Exception, r"Backend is not available."):
+        with self.assertRaisesRegexWithHighlight(Exception,
+                                                 r"Backend is not available.",
+                                                 "raise Exception(\"Backend is not available.\""):
             imported = torch.jit.load(buffer)
 
 
@@ -400,21 +404,23 @@
         Check errors associated with selective lowering.
         """
         # Check error messages thrown when attempting to lower something that is not a ScriptModule.
-        with self.assertRaisesRegex(RuntimeError, r"Object .* is not a ScriptModule"):
+        with self.assertRaisesRegexWithHighlight(RuntimeError, r"Object .* is not a ScriptModule", ""):
             to_test_backend_selective(torch.nn.ReLU(), {"forward": ""}, ["submodule"])
 
         MiddleModule = SelectiveLoweringTest.MiddleModule
         mod = MiddleModule(BasicModule())
         mod.new_attr = 3
 
-        with self.assertRaisesRegex(RuntimeError, r"Attribute named new_attr is not a Module"):
+        with self.assertRaisesRegexWithHighlight(RuntimeError, r"Attribute named new_attr is not a Module", ""):
             to_test_backend_selective(torch.jit.script(mod), {"forward": ""}, ["new_attr"])
 
         # Check error message thrown when module hierarchy doesn't have unique types.
         OuterModule = SelectiveLoweringTest.OuterModule
         mod = OuterModule(MiddleModule(BasicModule()), MiddleModule(BasicModule()), MiddleModule(BasicModule()))
 
-        with self.assertRaisesRegex(RuntimeError, r"Selective lowering is only supported for module hierarchies with unique types"):
+        with self.assertRaisesRegexWithHighlight(RuntimeError,
+                                                 r"Selective lowering is only supported for module hierarchies with unique types",
+                                                 ""):
             to_test_backend_selective(torch.jit.script(mod), {"forward": ""}, ["sub1.submodule"])