diff --git a/tests/unit/test_tensor.py b/tests/unit/test_tensor.py
index 7180e760..6395b444 100644
--- a/tests/unit/test_tensor.py
+++ b/tests/unit/test_tensor.py
@@ -55,7 +55,7 @@ def test_non_infer_returns(self):
             except NotImplementedError as e:
                 assert (
                     str(e)
-                    == "__mod__ with signature __mod__(self, other Any) -> Tensor"
+                    == "NotImplementedError: __mod__ with signature __mod__(self, other Any) -> Tensor"
                 )
 
     def test_optional_args(self):