search

shibing624 / textgen

TextGen: Implementation of Text Generation models, include LLaMA, BLOOM, GPT2, BART, T5, SongNet and so on. 文本生成模型，实现了包括LLaMA，ChatGLM，BLOOM，GPT2，Seq2Seq，BART，T5，UDA等模型的训练和预测，开箱即用。
Apache License 2.0
935 stars 109 forks source link

想要开启GPU上训练的话咋设置？ #10

Closed huyi1989 closed 2 years ago

huyi1989 commented 2 years ago

麻烦请问，开启训练后，模型是在CPU上跑的 数据量大的话有点慢 有没有办法在GPU上运行？

shibing624 commented 2 years ago

各模型默认都是在gpu上训练的。你check下：device = torch.device("cuda" if torch.cuda.is_available() else "cpu")的device输出，和gpu环境配置。