How to specify pytorch device?

1-800-BAD-CODE / punctuators

Package for inference for punctuation, true-casing, and sentence boundary detection

23 stars 2 forks source link

Closed ivoras closed 4 months ago

ivoras commented 4 months ago

Thanks for this wrapper! But, how can we specify the (CUDA) device to be used for inference?

1-800-BAD-CODE commented 4 months ago

If you update to the latest version from pypi (0.0.6), the punctuation models should use CUDA when available, or fall back on CPU if not.

I don't have a local GPU so I didn't test the CUDA functionality.

$ pip install punctuators --upgrade