diff --git a/dev-upgrade/axons/axons.yaml b/dev-upgrade/axons/axons.yaml
index ea5a548..143d168 100644
--- a/dev-upgrade/axons/axons.yaml
+++ b/dev-upgrade/axons/axons.yaml
@@ -27,15 +27,6 @@ metadata:
   labels:
     app: axons
 spec:
-  affinity:
-    nodeAffinity:
-      requiredDuringSchedulingIgnoredDuringExecution:
-        nodeSelectorTerms:
-        - matchExpressions:
-          - key: app.kubernetes.io/name
-            operator: In
-            values:
-            - trainer
   serviceName: axons
   replicas: 1
   selector:
@@ -46,6 +37,15 @@ spec:
       labels:
         app: axons
     spec:
+      affinity:
+        nodeAffinity:
+          requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+            - matchExpressions:
+              - key: app.kubernetes.io/name
+                operator: In
+                values:
+                - trainer
       containers:
       - name: axons
         image: axoniq/axonserver