diff --git a/pkg/controller.v1/tensorflow/pod.go b/pkg/controller.v1/tensorflow/pod.go
index bc232db577..c60539f341 100644
--- a/pkg/controller.v1/tensorflow/pod.go
+++ b/pkg/controller.v1/tensorflow/pod.go
@@ -269,10 +269,7 @@ func isDistributed(tfjob *tfv1.TFJob) bool {
 			}
 		}
 	}
-	if distributionCount != 1 {
-		return true
-	}
-	return false
+	return distributionCount != 1
 }
 
 func setRestartPolicy(podTemplateSpec *v1.PodTemplateSpec, spec *common.ReplicaSpec) {