1 files changed, 3 insertions, 4 deletions

diff --git a/training/functional.py b/training/functional.py
index 1fdfdc8..2da0f69 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -68,9 +68,8 @@ class TrainingStrategy():
 
 def get_models(
     pretrained_model_name_or_path: str,
-    emb_r: int = 8,
-    emb_lora_alpha: int = 8,
-    emb_lora_dropout: float = 0.0
+    emb_alpha: int = 8,
+    emb_dropout: float = 0.0
 ):
     tokenizer = MultiCLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder='tokenizer')
     text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder')
@@ -80,7 +79,7 @@ def get_models(
     sample_scheduler = UniPCMultistepScheduler.from_pretrained(
         pretrained_model_name_or_path, subfolder='scheduler')
 
-    embeddings = patch_managed_embeddings(text_encoder, emb_r, emb_lora_alpha, emb_lora_dropout)
+    embeddings = patch_managed_embeddings(text_encoder, emb_alpha, emb_dropout)
 
     return tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler, embeddings