Update README.md

README.md

@@ -16,8 +16,10 @@ pipeline_tag: text2text-generation
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
+device = "cuda"
+
 tokenizer = AutoTokenizer.from_pretrained('TeraSpace/dialofred')
-model = AutoModelForSeq2SeqLM.from_pretrained('TeraSpace/dialofred', device_map={'':0}).to(device)# Add torch_dtype=torch.bfloat16 to use less memory
+model = AutoModelForSeq2SeqLM.from_pretrained('TeraSpace/dialofred', device_map=device).to(device)# Add torch_dtype=torch.bfloat16 to use less memory
 while True:
     text_inp = input("=>")
     lm_text=f'<SC1>- {text_inp}\n- <extra_id_0>'