change to tf job worker spec

Signed-off-by: Future Outlier <[email protected]>
flyteorg · Oct 4, 2023 · e6aadd8 · e6aadd8
1 parent a9f5d24
commit e6aadd8
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(),