Update app.py

app.py

@@ -7,6 +7,7 @@ model_path = "ibm-granite/granite-vision-3.1-2b-preview"
 processor = LlavaNextProcessor.from_pretrained(model_path, use_fast=True)
 model = LlavaNextForConditionalGeneration.from_pretrained(model_path, torch_dtype="auto", device_map="auto")
 
+@spaces.GPU()
 def get_text_from_content(content):
     texts = []
     for item in content:
@@ -39,7 +40,7 @@ def chat_inference(image, text, temperature, top_p, top_k, max_tokens, conversat
         tokenize=True,
         return_dict=True,
         return_tensors="pt"
-    ).to("cpu")
+    ).to(model.device)
 
     torch.manual_seed(random.randint(0, 10000))