Fix variable reference

llms/mlx_lm/lora.py

@@ -235,6 +235,10 @@ def train_model(
             build_schedule(args.lr_schedule) if args.lr_schedule else args.learning_rate
         )
     )
+
+    if args.mask_inputs:
+        print("Masking inputs..")
+
     # Train model
     train(
         model=model,

llms/mlx_lm/tuner/trainer.py

@@ -168,7 +168,7 @@ def iterate_delineated_batches(
             for j in batch_idx[i]:
                 prompt, completion = dataset.get_prompt_and_completion(j)
                 prompt_length, completion_length = input_and_output_lengths(
-                    prompt, prompt, tokenizer
+                    prompt, completion, tokenizer
                 )
 
                 prompt_lengths.append(prompt_length)