addressed comments

kubeflow · Jul 26, 2019 · 6ae7ed3 · 6ae7ed3
1 parent 93a6a88
commit 6ae7ed3
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/sdk/python/kfp/compiler/_component_builder.py b/sdk/python/kfp/compiler/_component_builder.py
@@ -339,9 +339,11 @@ def _generate_entrypoint(self, component_func, python_version='python3'):
     func_signature = 'def ' + new_func_name + '('
     for input_arg in input_args:
       func_signature += input_arg + ','
-    func_signature += '_output_file):'
     if output_is_named_tuple:
-      func_signature = func_signature.replace('file', 'files')
+      func_signature += '_output_files'
+    else:
+      func_signature += '_output_file'
+    func_signature += '):'
     codegen.writeline(func_signature)
 
     # Call user function