diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py
index 642b1e8e27..5323e73ae1 100644
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@@ -143,7 +143,8 @@ class TrainingBackend:
             "dataset_slice_start": kwargs.get("dataset_slice_start"),
             "dataset_slice_end": kwargs.get("dataset_slice_end"),
             "custom_format_mapping": kwargs.get("custom_format_mapping"),
-            "is_dataset_multimodal": kwargs.get("is_dataset_multimodal", False),
+            "is_dataset_image": kwargs.get("is_dataset_image", False),
+            "is_dataset_audio": kwargs.get("is_dataset_audio", False),
             "num_epochs": kwargs.get("num_epochs", 3),
             "learning_rate": kwargs.get("learning_rate", "2e-4"),
             "batch_size": kwargs.get("batch_size", 2),