EricFillion / happy-transformer

Happy Transformer makes it easy to fine-tune and perform inference with NLP Transformer models.
http://happytransformer.com
Apache License 2.0
516 stars 66 forks source link

Is there an option to run on gpu? #339

Open Alaadel opened 3 months ago

Alaadel commented 3 months ago

Hello,

Whenever I run this code:

from happytransformer import HappyTextToText, TTSettings

happy_tt = HappyTextToText("T5", "vennify/t5-base-grammar-correction")

args = TTSettings(num_beams=5, min_length=1)

I get:

[timestamp] - INFO - happytransformer.happy_transformer -   Using device: cpu

Is there a way to force run on gpu? I have GTX 1060 6GB. Does it run automatically on gpu if certain conditions are met? Is this model-specific, or can be controlled by happy-transformer?

Thank you.