Update README.md
Browse files
README.md
CHANGED
@@ -19,7 +19,7 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
|
|
19 |
device = "cuda"
|
20 |
|
21 |
tokenizer = AutoTokenizer.from_pretrained('TeraSpace/dialofred')
|
22 |
-
model = AutoModelForSeq2SeqLM.from_pretrained('TeraSpace/dialofred', device_map=device)
|
23 |
while True:
|
24 |
text_inp = input("=>")
|
25 |
lm_text=f'<SC1>- {text_inp}\n- <extra_id_0>'
|
|
|
19 |
device = "cuda"
|
20 |
|
21 |
tokenizer = AutoTokenizer.from_pretrained('TeraSpace/dialofred')
|
22 |
+
model = AutoModelForSeq2SeqLM.from_pretrained('TeraSpace/dialofred', device_map=device)# Add torch_dtype=torch.bfloat16 to use less memory
|
23 |
while True:
|
24 |
text_inp = input("=>")
|
25 |
lm_text=f'<SC1>- {text_inp}\n- <extra_id_0>'
|