Update app.py

app.py

@@ -16,9 +16,7 @@ def get_model():
 
 def generate_text(prompt, max_tokens=500, temperature=0.3, top_k=40):
     """Generate text based on the prompt."""
-    # Encode the prompt
-    device = next(model.parameters()).device
-    
+       
     # Tokenize prompt with special token handling
     enc = tiktoken.get_encoding("gpt2")
     input_ids = torch.tensor(enc.encode(prompt, allowed_special={'<|endoftext|>'})).unsqueeze(0).to(device)
@@ -28,6 +26,7 @@ def generate_text(prompt, max_tokens=500, temperature=0.3, top_k=40):
     
     # Get cached model
     model = get_model()
+    device = next(model.parameters()).device
     
     with torch.no_grad():
         output_sequence = []