Update README.md
Browse files
README.md
CHANGED
|
@@ -95,7 +95,7 @@ print("Pipeline Output:", result)
|
|
| 95 |
import torch
|
| 96 |
|
| 97 |
# Function for direct tokenization and text generation
|
| 98 |
-
def generate_text(input_text, device):
|
| 99 |
tokens = tokenizer.encode(input_text, return_tensors='pt').to(device)
|
| 100 |
model.to(device)
|
| 101 |
|
|
|
|
| 95 |
import torch
|
| 96 |
|
| 97 |
# Function for direct tokenization and text generation
|
| 98 |
+
def generate_text(input_text, device='cpu'):
|
| 99 |
tokens = tokenizer.encode(input_text, return_tensors='pt').to(device)
|
| 100 |
model.to(device)
|
| 101 |
|