change to tf job worker spec (#4169)

Signed-off-by: Future Outlier <[email protected]> Co-authored-by: Future Outlier <[email protected]>
flyteorg · Oct 5, 2023 · c88f8f5 · c88f8f5
1 parent 81b94a4
commit c88f8f5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go
@@ -122,7 +122,7 @@ func (tensorflowOperatorResourceHandler) BuildResource(ctx context.Context, task
 		workerReplicaSpec := kfTensorflowTaskExtraArgs.GetWorkerReplicas()
 		if workerReplicaSpec != nil {
 			err := common.OverrideContainerSpec(
-				replicaSpecMap[kubeflowv1.MPIJobReplicaTypeWorker].PodSpec,
+				replicaSpecMap[kubeflowv1.TFJobReplicaTypeWorker].PodSpec,
 				kubeflowv1.TFJobDefaultContainerName,
 				workerReplicaSpec.GetImage(),
 				workerReplicaSpec.GetResources(),