diff --git a/README.md b/README.md
index 47355f8..b800f42 100644
--- a/README.md
+++ b/README.md
@@ -47,12 +47,13 @@ checkpoint = "bigcode/starcoder"
 device = "cuda" # for GPU usage or "cpu" for CPU usage
 
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-# to save memory consider using fp16 or bf16 by specifying torch.dtype=torch.float16 for example
+# to save memory consider using fp16 or bf16 by specifying torch_dtype=torch.float16 for example
 model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
 
 inputs = tokenizer.encode("def print_hello_world():", return_tensors="pt").to(device)
 outputs = model.generate(inputs)
-print(tokenizer.decode(outputs[0]))
+# clean_up_tokenization_spaces=False prevents a tokenizer edge case which can result in spaces being removed around punctuation
+print(tokenizer.decode(outputs[0], clean_up_tokenization_spaces=False))
 ```
 or
 ```python
diff --git a/chat/generate.py b/chat/generate.py
index 64a3905..c67d7d8 100644
--- a/chat/generate.py
+++ b/chat/generate.py
@@ -127,7 +127,7 @@ def main():
         print()
 
     raw_model_name = args.model_id.split("/")[-1]
-    model_name = f"{raw_model_name}-{args.prompt_type}"
+    model_name = f"{raw_model_name}"
     if args.revision is not None:
         model_name += f"-{args.revision}"
 
diff --git a/finetune/finetune.py b/finetune/finetune.py
index 96ab961..525b37f 100644
--- a/finetune/finetune.py
+++ b/finetune/finetune.py
@@ -267,6 +267,8 @@ def run_training(args, train_data, val_data):
         output_dir=args.output_dir,
         dataloader_drop_last=True,
         evaluation_strategy="steps",
+        save_strategy="steps",
+        load_best_model_at_end=True,
         max_steps=args.max_steps,
         eval_steps=args.eval_freq,
         save_steps=args.save_freq,
@@ -309,4 +311,4 @@ def main(args):
 
     logging.set_verbosity_error()
 
-    main(args)
\ No newline at end of file
+    main(args)
diff --git a/finetune/merge_peft_adapters.py b/finetune/merge_peft_adapters.py
index 790b6fd..218728f 100644
--- a/finetune/merge_peft_adapters.py
+++ b/finetune/merge_peft_adapters.py
@@ -25,7 +25,6 @@ def main():
     model = PeftModel.from_pretrained(base_model, args.peft_model_path)
     model = model.merge_and_unload()
 
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     tokenizer = AutoTokenizer.from_pretrained(args.base_model_name_or_path)
 
     if args.push_to_hub:
@@ -38,4 +37,4 @@ def main():
         print(f"Model saved to {args.base_model_name_or_path}-merged")
 
 if __name__ == "__main__" :
-    main()
\ No newline at end of file
+    main()