diff --git a/tensorflow/python/client/tf_session.i b/tensorflow/python/client/tf_session.i
index 8e8e3132f1a..05ffcc93504 100644
--- a/tensorflow/python/client/tf_session.i
+++ b/tensorflow/python/client/tf_session.i
@@ -208,7 +208,12 @@ tensorflow::ImportNumpy();
     RaiseStatusNotOK(*$1, $descriptor(tensorflow::Status*));
     SWIG_fail;
   } else {
-    $result = PyString_FromStringAndSize(*$2, strlen(*$2));
+%#if PY_MAJOR_VERSION < 3
+    $result = PyString_FromStringAndSize(
+%#else
+    $result = PyUnicode_FromStringAndSize(
+%#endif
+      *$2, strlen(*$2));
     delete *$2;
   }
 }