Update app.py

app.py

@@ -257,23 +257,22 @@ def process_images(images, threshold):
 
     with torch.no_grad():
         for batch, filenames in dataloader:
-            
             batch = batch.to(device) 
-            with torch.no_grad():
-                probabilities = model(batch)
-                for i, prob in enumerate(probabilities):
-                    prob = prob[0]
-                    indices = torch.where(prob > threshold)[0]
-                    values = prob[indices]
-                    
-                    temp = []
-                    tag_score = dict()
-                    for j in range(indices.size(0)):
-                        temp.append([allowed_tags[indices[j]], values[j].item()])
-                        tag_score[allowed_tags[indices[j]]] = values[j].item()
-                    
-                    tags = ", ".join([t[0] for t in temp])
-                    all_results.append((filenames[i], tags, tag_score))
+            probabilities = model(batch)
+            for i, prob in enumerate(probabilities):
+                indices = torch.where(prob > threshold)[0]
+                values = prob[indices]
+                
+                temp = []
+                tag_score = dict()
+                for j in range(indices.size(0)):
+                    tag = allowed_tags[indices[j]]
+                    score = values[j].item()
+                    temp.append([tag, score])
+                    tag_score[tag] = score
+                
+                tags = ", ".join([t[0] for t in temp])
+                all_results.append((filenames[i], tags, tag_score))
 
     return all_results