diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go
index c8beacc950..578bd2a0d6 100644
--- 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(),