diff --git a/.helm/templates/sts.yaml b/.helm/templates/sts.yaml
index a108dfa..60c664c 100644
--- a/.helm/templates/sts.yaml
+++ b/.helm/templates/sts.yaml
@@ -1,6 +1,3 @@
-{{- $gpuRole := .Values.gpu.nodeRoleKey }}
-{{- $tolKey := .Values.gpu.tolerationKey }}
-{{- $tolVal := .Values.gpu.tolerationValue }}
 ---
 apiVersion: apps/v1
 kind: StatefulSet
@@ -31,20 +28,11 @@ spec:
                   app: {{ .Chart.Name }}
               topologyKey: kubernetes.io/hostname
             weight: 100
-        {{- if $gpuRole }}
-        nodeAffinity:
-          requiredDuringSchedulingIgnoredDuringExecution:
-            nodeSelectorTerms:
-            - matchExpressions:
-              - key: {{ $gpuRole | quote }}
-                operator: Exists
-        {{- end }}
-      {{- if and $tolKey $tolVal }}
-      tolerations:
-      - key: {{ $tolKey | quote }}
-        operator: Equal
-        value: {{ $tolVal | quote }}
-        effect: NoExecute
+      {{- if and .Values.gpu.enabled .Values.gpu.nodeSelector }}
+      nodeSelector: {{ toYaml .Values.gpu.nodeSelector | nindent 8 }}
+      {{- end }}
+      {{- if and .Values.gpu.enabled .Values.gpu.tolerations }}
+      tolerations: {{ toYaml .Values.gpu.tolerations | nindent 6 }}
       {{- end }}
       imagePullSecrets:
       - name: registrysecret
diff --git a/.helm/values.yaml b/.helm/values.yaml
index 139db74..23bffbc 100644
--- a/.helm/values.yaml
+++ b/.helm/values.yaml
@@ -17,9 +17,14 @@ env:
   TRANSFORMERS_OFFLINE: 1
 
 gpu:
-  nodeRoleKey: node-role.deckhouse.io/worker
-  tolerationKey: dedicated
-  tolerationValue: ""
+  enabled: true
+  nodeSelector:
+    node-role.deckhouse.io/cuda-worker: ""
+  tolerations:
+  - key: dedicated.deckhouse.io
+    operator: Equal
+    value: cuda-worker
+    effect: NoExecute
 
 resources:
   _default: