Update handler.py

handler.py

@@ -181,7 +181,7 @@ class EndpointHandler:
             outputs = self.safety_checker(**inputs)
             logits_per_image = outputs.logits_per_image
             probs = logits_per_image.softmax(dim=1)  # We assume the probability for NSFW content is stored in the first position
-            return probs[0, 0] > 0.5  # This threshold may need to be adjusted
+            return probs[0, 1] > 0.5  # This threshold may need to be adjusted
 
         request = GenerateImageRequest(**data)
         inputs = request.inputs