Update app.py

app.py

@@ -12,7 +12,10 @@ class GPTConfig:
     n_head: int = 12 # number of heads
     n_embd: int = 768 # embedding dimension
 
-
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model2 = model.GPT(GPTConfig())
+model2 = model2.to(device)
+model2.load_state_dict(torch.load('gpt_124M_1.pth', map_location=torch.device(device)))
 
 # Define generation function
 def generate_text(prompt, max_length=50, num_return_sequences=10):