Update README.md

README.md

@@ -22,12 +22,16 @@ GGUF
 ## Usage
 
 ~~~Bash
-pip install auto-gptq[triton]==0.4.2 transformers
+pip install auto-gptq[triton]==0.4.2 transformers==4.34.1
 ~~~
   
 ~~~python
+import torch
 from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer , AutoModelForCausalLM
+
+if torch.cuda.is_available():
+    device_name = torch.cuda.get_device_name(0)
 
 model_name_or_path = "mmnga/cyberagent-calm2-7b-chat-GPTQ-calib-ja-1k"
 
@@ -35,7 +39,7 @@ model_name_or_path = "mmnga/cyberagent-calm2-7b-chat-GPTQ-calib-ja-1k"
 tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
 
 # Model
-model = AutoGPTQForCausalLM.from_quantized(model_name_or_path, use_safetensors=True, device="cuda:0")
+model = AutoGPTQForCausalLM.from_quantized(model_name_or_path, use_safetensors=True, device="cuda:0", use_triton=("A100" in device_name))
 
 # Your test prompt
 prompt = """