Update README.md

README.md

@@ -28,7 +28,7 @@ import torch
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 device = "cuda" # you can use "auto" for placing the model on several GPUs if your GPUs can not fit the model
 path = "nicolinho/QRM-Gemma-2-27B"
-model = AutoModelForSequenceClassification.from_pretrained(path, torch_dtype=torch.bfloat16, device_map=device, trust_remote_code=True)
+model = AutoModelForSequenceClassification.from_pretrained(path, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", device_map=device, trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained(path, use_fast=True)
 # We load a random sample from the validation set of the HelpSteer dataset
 prompt = 'Does pineapple belong on a Pizza?'