Update app.py

app.py

@@ -9,7 +9,10 @@ import requests
 from samplings import top_p_sampling, top_k_sampling, temperature_sampling
 from transformers import GPT2Config, GPT2Model, GPT2LMHeadModel, PreTrainedModel
 
-device = torch.device("cpu")
+if torch.cuda.is_available():    
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
 
 description = """
 <div>
@@ -130,7 +133,7 @@ def generate_abc(prompt,
             print(f"Error: {e}")
             exit()
             
-    model.load_state_dict(torch.load('pytorch_model.bin'))
+    model.load_state_dict(torch.load('pytorch_model.bin', map_location=device))
     model.eval()
 
     tunes = ""