diff --git a/docker/helm/templates/master-statefulset.yaml b/docker/helm/templates/master-statefulset.yaml
index f47d38ecd..90009cec2 100644
--- a/docker/helm/templates/master-statefulset.yaml
+++ b/docker/helm/templates/master-statefulset.yaml
@@ -42,7 +42,10 @@ spec:
         app.kubernetes.io/role: master
         {{- include "celeborn.selectorLabels" . | nindent 8 }}
     spec:
+      {{- with .Values.affinity }}
       affinity:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
         podAntiAffinity:
           requiredDuringSchedulingIgnoredDuringExecution:
           - labelSelector:
@@ -60,10 +63,6 @@ spec:
       nodeSelector:
         {{- toYaml . | nindent 8 }}
       {{- end }}
-      {{- with .Values.affinity }}
-      affinity:
-        {{- toYaml . | nindent 8 }}
-      {{- end }}
       {{- with .Values.tolerations }}
       tolerations:
         {{- toYaml . | nindent 8 }}
diff --git a/docker/helm/templates/worker-statefulset.yaml b/docker/helm/templates/worker-statefulset.yaml
index 196f01f81..4fdf6e9dd 100644
--- a/docker/helm/templates/worker-statefulset.yaml
+++ b/docker/helm/templates/worker-statefulset.yaml
@@ -42,7 +42,10 @@ spec:
         app.kubernetes.io/role: worker
         {{- include "celeborn.selectorLabels" . | nindent 8 }}
     spec:
+      {{- with .Values.affinity }}
       affinity:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
         podAntiAffinity:
           requiredDuringSchedulingIgnoredDuringExecution:
           - labelSelector:
@@ -60,10 +63,6 @@ spec:
       nodeSelector:
         {{- toYaml . | nindent 8 }}
       {{- end }}
-      {{- with .Values.affinity }}
-      affinity:
-        {{- toYaml . | nindent 8 }}
-      {{- end }}
       {{- with .Values.tolerations }}
       tolerations:
         {{- toYaml . | nindent 8 }}