Override requests and limits in kwargs with values from task_config (f…

…lyteorg#419) As done in kftensorflow in flyteorg#408 Signed-off-by: Sören Brunk <[email protected]>
dolthub · Mar 14, 2021 · 805c5c2 · 805c5c2
1 parent 7600429
commit 805c5c2
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/plugins/kfpytorch/flytekitplugins/kfpytorch/task.py b/plugins/kfpytorch/flytekitplugins/kfpytorch/task.py
@@ -47,9 +47,7 @@ def __init__(self, task_config: PyTorch, task_function: Callable, **kwargs):
             task_config,
             task_function,
             task_type=self._PYTORCH_TASK_TYPE,
-            requests=task_config.per_replica_requests,
-            limits=task_config.per_replica_limits,
-            **kwargs
+            **{**kwargs, "requests": task_config.per_replica_requests, "limits": task_config.per_replica_limits}
         )
 
     def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]: