diff --git a/webapp/src/components/training/TrainingSteps.vue b/webapp/src/components/training/TrainingSteps.vue
index 7f9fc5c52..f4130e24b 100644
--- a/webapp/src/components/training/TrainingSteps.vue
+++ b/webapp/src/components/training/TrainingSteps.vue
@@ -1,17 +1,16 @@
 <template>
-  <div v-if="task !== undefined">
+  <div v-if="task !== undefined" class="flex flex-col space-y-4 md:space-y-8">
     <Description v-show="trainingStore.step === 1" :task="task" />
 
-    <div
+    <LabeledDatasetInput
       v-show="trainingStore.step === 2"
-      class="flex flex-col space-y-4 md:space-y-8"
+      :task
+      v-model="dataset"
     >
-      <LabeledDatasetInput :task v-model="dataset">
-        <template #header>
-          <DataDescription :task />
-        </template>
-      </LabeledDatasetInput>
-    </div>
+      <template #header>
+        <DataDescription :task />
+      </template>
+    </LabeledDatasetInput>
 
     <Trainer
       v-show="trainingStore.step === 3"
@@ -21,8 +20,9 @@
     />
 
     <Finished v-show="trainingStore.step === 4" :task :model="trainedModel" />
+
+    <TrainingButtons />
   </div>
-  <TrainingButtons />
 </template>
 
 <script lang="ts" setup>