Update README.md

README.md

@@ -221,23 +221,13 @@ RUSpellRU, MultidomainGold, MedSpellChecker, GitHubTypoCorpusRu are datasets for
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 tokenizer = AutoTokenizer.from_pretrained("ai-forever/sage-mt5-large")
-model = AutoModelForSeq2SeqLM.from_pretrained("ai-forever/sage-mt5-large")
-model.to("cuda:0")
+model = AutoModelForSeq2SeqLM.from_pretrained("ai-forever/sage-mt5-large", device_map='cuda')
 
 sentence = "Перведи мне текст на аглиском: \"Screw you kuys, I am goin hme (c)."
-with torch.inference_mode():
-    encodings = tokenizer(sentence, max_length=None, padding="longest", truncation=False, return_tensors="pt")
-    for k, v in encodings.items():
-        encodings[k] = v.to("cuda:0")
-    res = model.generate(
-        **encodings,
-        use_cache=True,
-        max_length = encodings["input_ids"].size(1) * 1.5
-    )
-    res = res.cpu().tolist()
-    res = tokenizer.batch_decode(res, skip_special_tokens=True)
+inputs = tokenizer(sentence, max_length=None, padding="longest", truncation=False, return_tensors="pt")
+outputs = model.generate(**inputs.to(model.device), max_length = inputs["input_ids"].size(1) * 1.5)
+print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
 
-print(res)
 # ["Переведи мне текст на английском: "Screw you guys, I am going home" (c)."]
 ```