Update README.md
Browse files
README.md
CHANGED
|
@@ -50,6 +50,9 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
|
|
| 50 |
|
| 51 |
models_folder = "Deeokay/gpt2-javis-stks"
|
| 52 |
|
|
|
|
|
|
|
|
|
|
| 53 |
model = GPT2LMHeadModel.from_pretrained(models_folder)
|
| 54 |
tokenizer = GPT2Tokenizer.from_pretrained(models_folder)
|
| 55 |
|
|
|
|
| 50 |
|
| 51 |
models_folder = "Deeokay/gpt2-javis-stks"
|
| 52 |
|
| 53 |
+
# if you know your device, you can just set "device = 'mps'"
|
| 54 |
+
device = torch.device("cuda" if torch.cuda.is_available90 else "cpu")
|
| 55 |
+
|
| 56 |
model = GPT2LMHeadModel.from_pretrained(models_folder)
|
| 57 |
tokenizer = GPT2Tokenizer.from_pretrained(models_folder)
|
| 58 |
|