Moved model to gpu before calling count functions

app.py

@@ -222,6 +222,8 @@ def get_ind_to_filter(text, word_ids, keywords):
 
 @spaces.GPU
 def count(image, text, prompts, state, device):
+
+    model.to(device)
     
     keywords = "" # do not handle this for now
     
@@ -314,6 +316,8 @@ def count(image, text, prompts, state, device):
 
 @spaces.GPU
 def count_main(image, text, prompts, device):
+    model.to(device)
+    
     keywords = "" # do not handle this for now
     # Handle no prompt case.
     if prompts is None: