huggingface / huggingface/evaluate

Toxicity Measurement Should Allow GPU Usage

Open
#454 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
2.5k
Forks
341
PR merge metrics
No merged PRs in 30d

Description

The toxicity measurement is based on BERT models in HuggingFace, and thus, should be able to utilize GPUs at inference time if a user wants. Suggesting to add the following line to `_download_and_prepare`:

`device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")`

And then add `device=device` param to `pipeline`.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.