use gpu if possible

app/app.py

@@ -16,6 +16,7 @@ mirror_url = "https://news-generator.ai-research.id/"
 if "MIRROR_URL" in os.environ:
     mirror_url = os.environ["MIRROR_URL"]
 hf_auth_token = os.getenv("HF_AUTH_TOKEN", False)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 MODELS = {
     "Indonesian Newspaper - Indonesian GPT-2 Medium": {
@@ -63,6 +64,7 @@ def get_generator(model_name: str):
     st.write(f"Loading the GPT2 model {model_name}, please wait...")
     tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_auth_token)
     model = GPT2LMHeadModel.from_pretrained(model_name, pad_token_id=tokenizer.eos_token_id, use_auth_token=hf_auth_token)
+    model.to(device)
     model.resize_token_embeddings(len(tokenizer))
     return model, tokenizer
 
@@ -82,8 +84,7 @@ def process(text_generator, tokenizer, title: str, keywords: str, text: str,
     prompt = f"title: {title}\nkeywords: {keywords}\n{text}"
 
     generated = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)
-    # device = torch.device("cuda")
-    # generated = generated.to(device)
+    generated = generated.to(device)
 
     text_generator.eval()
     sample_outputs = text_generator.generate(generated,