fix dreambooth colab

huggingface · May 22, 2024 · 30e757b · 30e757b
1 parent 16c20bc
commit 30e757b
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/src/autotrain/app/colab.py b/src/autotrain/app/colab.py
@@ -351,7 +351,7 @@ def start_training(b):
                     "backend": "local",
                     "data": {
                         "path": dataset_path.value,
-                        "prompt": params_val["prompt"],
+                        "prompt": prompt,
                     },
                     "params": params_val,
                     "hub": {

diff --git a/src/autotrain/commands.py b/src/autotrain/commands.py
@@ -17,6 +17,18 @@
 
 
 def launch_command(params):
+    """
+    Launches training command based on the given parameters.
+
+    Args:
+        params: An instance of a parameter class (LLMTrainingParams, DreamBoothTrainingParams, GenericParams, TabularParams,
+                TextClassificationParams, TextRegressionParams, TokenClassificationParams, ImageClassificationParams,
+                ObjectDetectionParams, Seq2SeqParams).
+
+    Returns:
+        None
+    """
+
     params.project_name = shlex.split(params.project_name)[0]
     cuda_available = torch.cuda.is_available()
     mps_available = torch.backends.mps.is_available()