Skip to content

Commit 8562a22

Browse files
author
Roja Reddy Sareddy
committed
fix: Fix the recipe selection for multiple recipe scenario
1 parent c4e9ace commit 8562a22

File tree

1 file changed

+2
-4
lines changed

1 file changed

+2
-4
lines changed

sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -347,11 +347,9 @@ def _get_fine_tuning_options_and_model_arn(model_name: str, customization_techni
347347

348348
# Select recipe based on training type
349349
recipe = None
350-
if isinstance(training_type, TrainingType) and training_type == TrainingType.LORA:
351-
# For LORA, find first recipe with Peft key
350+
if (isinstance(training_type, TrainingType) and training_type == TrainingType.LORA) or training_type == "LORA":
352351
recipe = next((r for r in recipes_with_template if r.get("Peft")), None)
353-
elif isinstance(training_type, TrainingType) and training_type == TrainingType.FULL:
354-
# For FULL, find first recipe without Peft key
352+
elif (isinstance(training_type, TrainingType) and training_type == TrainingType.FULL) or training_type == "FULL":
355353
recipe = next((r for r in recipes_with_template if not r.get("Peft")), None)
356354

357355
if recipe and recipe.get("SmtjOverrideParamsS3Uri"):

0 commit comments

Comments
 (0)