Update README.md

README.md

@@ -34,7 +34,7 @@ torch.manual_seed(42)
 # script config
 base_model_name = "LSX-UniWue/LLaMmlein_1B"
 chat_adapter_name = "LSX-UniWue/LLaMmlein_1B_chat_all"
-device = "mps"  # or cuda
+device = "cuda"  # or mps
 
 # chat history
 messages = [
@@ -48,7 +48,6 @@ messages = [
 config = PeftConfig.from_pretrained(chat_adapter_name)
 base_model = model = AutoModelForCausalLM.from_pretrained(
     base_model_name,
-    attn_implementation="flash_attention_2" if device == "cuda" else None,
     torch_dtype=torch.bfloat16,
     device_map=device,
 )