Update app.py

app.py

@@ -3,13 +3,16 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import spaces
 
+zero = torch.Tensor([0]).cuda()
+
 # Load model and tokenizer only once, outside the function
 model_name = "deepapaikar/Katzbot_Llama_7b_QA_10eps"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name,  device_map='auto')
 
 
-@spaces
+
+@spaces.GPU
 def generate_text(input_text):
     """Generates text using the LlamaKatz-3x8B model.
 
@@ -19,7 +22,7 @@ def generate_text(input_text):
     Returns:
         str: The generated text.
     """
-    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+    inputs = tokenizer(input_text, return_tensors="pt").to(zero.device)
     outputs = model.generate(**inputs)
     generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return generated_text