Update app.py

app.py

@@ -141,7 +141,8 @@ def run_classifier(image: Image.Image, threshold):
     tensor = transform(img).unsqueeze(0)
 
     with torch.no_grad():
-        probits = model(tensor)[0] # type: torch.Tensor
+        logits = model(tensor)
+        probits = torch.nn.functional.sigmoid(logits[0])
         values, indices = probits.cpu().topk(250)
 
     tag_score = {allowed_tags[idx.item()]: val.item() for idx, val in zip(indices, values)}
@@ -175,7 +176,8 @@ def cam_inference(img, threshold, alpha, evt: gr.SelectData):
     handle_forward = model.norm.register_forward_hook(hook_forward)
     handle_backward = model.norm.register_full_backward_hook(hook_backward)
 
-    probits = model(tensor)[0]
+    logits = model(tensor)
+    probits = torch.nn.functional.sigmoid(logits[0])
  
     model.zero_grad()
     probits[target_tag_index].backward(retain_graph=True)