added @spaces.GPU

app.py

@@ -1,6 +1,12 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
+import spaces
+import torch
+
+
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+print(f'[INFO] Using device: {device}')
 
 # token
 token = os.environ['TOKEN']
@@ -9,7 +15,7 @@ token = os.environ['TOKEN']
 MODEL_NAME = "atlasia/Al-Atlas-LLM"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=token)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, token=token).to('cuda')
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, token=token).to(device)
 
 # Predefined examples
 examples = [
@@ -23,6 +29,7 @@ examples = [
      , 256, 0.7, 0.9, 150, 8, 1.5],
 ]
 
+@spaces.GPU
 def generate_text(prompt, max_length=256, temperature=0.7, top_p=0.9, top_k=150, num_beams=8, repetition_penalty=1.5):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     output = model.generate(