Update app.py

app.py

@@ -315,7 +315,7 @@ def count(image, text, prompts, state, device):
 @spaces.GPU
 def count_main(image, text, prompts, device):
     model.to(device)
-    print("model device: " + str(model.device))
+    
     keywords = "" # do not handle this for now
     # Handle no prompt case.
     if prompts is None:
@@ -327,6 +327,7 @@ def count_main(image, text, prompts, device):
     input_image_exemplars, exemplars = transform(prompts["image"], {"exemplars": torch.tensor(exemplars)})
     input_image_exemplars = input_image_exemplars.unsqueeze(0).to(device)
     exemplars = [exemplars["exemplars"].to(device)]
+    print("image device: " + str(input_image.device))
     
     with torch.no_grad():
         model_output = model(