Returned to GPU

app.py

@@ -112,7 +112,7 @@ def get_args_parser():
     parser.add_argument("--amp", action="store_true", help="Train with mixed precision")
     return parser
 
-#@spaces.GPU
+@spaces.GPU
 # Get counting model.
 def build_model_and_transforms(args):
     normalize = T.Compose(
@@ -219,7 +219,7 @@ def get_ind_to_filter(text, word_ids, keywords):
 
     return inds_to_filter
 
-#@spaces.GPU
+@spaces.GPU
 def count(image, text, prompts, state, device):
     print("state: " + str(state))
     keywords = "" # do not handle this for now
@@ -308,7 +308,7 @@ def count(image, text, prompts, state, device):
         out_label = "Nothing specified to detect."
     return (gr.Image(output_img, visible=True, label=out_label, show_label=True), gr.Number(label="Predicted Count", visible=True, value=boxes.shape[0]), new_submit_btn, gr.Tab(visible=True), step_3, state)
 
-#@spaces.GPU
+@spaces.GPU
 def count_main(image, text, prompts, device):
     keywords = "" # do not handle this for now
     # Handle no prompt case.