Update app.py

app.py

@@ -13,7 +13,6 @@ def infer(image, policies, policy_descriptions):
   policies = policies.split(";")
   policy_descriptions = policy_descriptions.split(";")
   custom_policies = dict(zip(policies, policy_descriptions))
-  print(custom_policies)
 
   inputs = processor(
       images=[image],
@@ -24,8 +23,8 @@ def infer(image, policies, policy_descriptions):
   
   with torch.inference_mode():
       output = model(**inputs)
-      print(output.probabilities)
-  
+      
+  output = output.cpu()
   outs = {}
   for idx, policy in enumerate(output.probabilities):
     yes_prob = policy[0]
@@ -33,7 +32,6 @@ def infer(image, policies, policy_descriptions):
 
     outs[f"Yes for {policies[idx]}"] = yes_prob
     outs[f"No for {policies[idx]}"] = no_prob
-    print(outs)
 
   return outs