Fix tensor printing bug in Python 2 (#12732)

Summary:
`rsplit` doesn't have kwargs in Python 2 so this line raises an error

Fixes #15135
Pull Request resolved: https://github.com/pytorch/pytorch/pull/12732

Differential Revision: D10458630

Pulled By: driazati

fbshipit-source-id: a63e42fbc0e39e4291480775b516c98122ec05a1

diff --git a/test/test_jit.py b/test/test_jit.py
index 81e8e5ae479811ca898f011bad3319d6c4d61971..77df56bc67c04bc1cc192427e2ab71b800d0b7ca 100644 (file)
--- a/test/test_jit.py
+++ b/test/test_jit.py
@@ -8764,6 +8764,18 @@ a")
 
         self.assertExpectedGraph(foo.graph)
 
+    def test_cpp_function_tensor_str(self):
+        x = torch.randn(2, 2)
+        scale = torch.randn(2, 2, requires_grad=True)
+        shift = torch.randn(2, 2, requires_grad=True)
+
+        @torch.jit.script
+        def fn(x, scale, shift):
+            return scale * x + shift
+
+        with self.capture_stdout() as captured:
+            print(fn(x, scale, shift))
+
 
 class MnistNet(nn.Module):
     def __init__(self):

diff --git a/torch/_tensor_str.py b/torch/_tensor_str.py
index a00c32e572e50c5b81cd871b9d95fecf8190b055..4132729b41e682ca9047f677cdf9281917c91521 100644 (file)
--- a/torch/_tensor_str.py
+++ b/torch/_tensor_str.py
@@ -282,7 +282,7 @@ def _str(self):
     if self.grad_fn is not None:
         name = type(self.grad_fn).__name__
         if name == 'CppFunction':
-            name = self.grad_fn.name().rsplit('::', maxsplit=1)[-1]
+            name = self.grad_fn.name().rsplit('::', 1)[-1]
         suffixes.append('grad_fn=<{}>'.format(name))
     elif self.requires_grad:
         suffixes.append('requires_grad=True')