Update app.py

app.py

@@ -24,9 +24,10 @@ def infer(image, policies, policy_descriptions):
   with torch.inference_mode():
       output = model(**inputs)
       
-  output = output.cpu()
+  
   outs = {}
-  for idx, policy in enumerate(output.probabilities):
+
+  for idx, policy in enumerate(output.probabilities.cpu()):
     yes_prob = policy[0]
     no_prob = policy[1]