Update app.py

app.py

@@ -260,8 +260,9 @@ def process_images(images, threshold):
             
             batch = batch.to(device) 
             with torch.no_grad():
-                probabilities = model(batch)[0]
+                probabilities = model(batch)
                 for i, prob in enumerate(probabilities):
+                    prob = prob[0]
                     indices = torch.where(prob > threshold)[0]
                     values = prob[indices]