change: minor changes based on feedback

Author: caxiaohu

src/sagemaker/tensorflow/estimator.py

@@ -531,15 +531,26 @@ def transformer(self, instance_count, instance_type, strategy=None, assemble_wit
                 If None, server will use one worker per vCPU.
             volume_kms_key (str): Optional. KMS key ID for encrypting the volume attached to the ML
                 compute instance (default: None).
-            endpoint_type: Optional. Selects the software stack used by the inference server.
-                If  not specified, the model will be configured to use the default
-                SageMaker model server. If 'tensorflow-serving', the model will be configured to
+            endpoint_type (str): Optional. Selects the software stack used by the inference server.
+                If not specified, the model will be configured to use the default
+                SageMaker model server.
+                If 'tensorflow-serving', the model will be configured to
                 use the SageMaker Tensorflow Serving container.
         """
 
         if endpoint_type == 'tensorflow-serving':
             self.script_mode = True
-        return super(TensorFlow, self).transformer(instance_count, instance_type, strategy, assemble_with, output_path,
-                                            output_kms_key, accept, env, max_concurrent_transforms, max_payload,
-                                            tags, role, model_server_workers, volume_kms_key
-                                            )
+        return super(TensorFlow, self).transformer(instance_count,
+                                                   instance_type,
+                                                   strategy,
+                                                   assemble_with,
+                                                   output_path,
+                                                   output_kms_key,
+                                                   accept,
+                                                   env,
+                                                   max_concurrent_transforms,
+                                                   max_payload,
+                                                   tags,
+                                                   role,
+                                                   model_server_workers,
+                                                   volume_kms_key)

tests/unit/test_tf_estimator.py

@@ -19,11 +19,11 @@
 import pytest
 from mock import patch, Mock, MagicMock
 
+from sagemaker.estimator import _TrainingJob
 from sagemaker.fw_utils import create_image_uri
 from sagemaker.model import MODEL_SERVER_WORKERS_PARAM_NAME
 from sagemaker.session import s3_input
 from sagemaker.tensorflow import defaults, TensorFlow, TensorFlowModel, TensorFlowPredictor
-from sagemaker.estimator import _TrainingJob
 import sagemaker.tensorflow.estimator as tfe
 from sagemaker.transformer import Transformer
 
@@ -293,6 +293,7 @@ def test_transformer_creation_with_endpoint_type(create_tfs_model, sagemaker_ses
     assert tf.script_mode is True
     assert tf._script_mode_enabled() is True
 
+
 @patch('sagemaker.tensorflow.estimator.TensorFlow._create_default_model')
 def test_transformer_creation_without_endpoint_type(create_default_model, sagemaker_session):
     container_log_level = '"logging.INFO"'