Skip to content

Instantly share code, notes, and snippets.

@charlieoneill11
Last active May 1, 2022 06:06
Show Gist options
  • Save charlieoneill11/2957b6a51b46cc999b9ce7024461ec49 to your computer and use it in GitHub Desktop.
Save charlieoneill11/2957b6a51b46cc999b9ce7024461ec49 to your computer and use it in GitHub Desktop.
from transformers import AutoModelForSequenceClassification
num_labels = 2
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = (AutoModelForSequenceClassification
.from_pretrained(model_ckpt, num_labels=num_labels)
.to(device))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment