AIP-3158: It actually needs to return a list of templates as it exten…

…ds the returned values
kubeflow · Jan 5, 2021 · 576f10c · 576f10c
1 parent 6df7887
commit 576f10c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/sdk/python/kfp/dsl/_container_op.py b/sdk/python/kfp/dsl/_container_op.py
@@ -927,11 +927,11 @@ def set_display_name(self, name: str):
         self.display_name = name
         return self
 
-    def set_template_handler(self, template_handler: Callable[['BaseOp'], Dict]):
+    def set_template_handler(self, template_handler: Callable[['BaseOp'], List[Dict]]):
         """Set the template handler for converting Op into Argo template.
 
         Args:
-          template_handler: A callable for converting the BaseOp into an Argo template during
+          template_handler: A callable for converting the BaseOp into Argo templates during
             pipeline compilation.
         """
         self.template_handler = template_handler