huggingface使用GPU推理

最新推荐文章于 2024-08-05 19:58:52 发布

hehui0921

最新推荐文章于 2024-08-05 19:58:52 发布

阅读量1k

点赞数 3

分类专栏： GPU 文章标签：人工智能

本文链接：https://blog.csdn.net/oHeHui1/article/details/136352660

版权

GPU 专栏收录该内容

1 篇文章 0 订阅

订阅专栏

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small",device_map="cuda:0")
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")

inputs = tokenizer("A step by step recipe to make bolognese pasta:", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))