pengdaqian commited on
Commit
7424ea1
1 Parent(s): 7f37981
Files changed (1) hide show
  1. img_nsfw.py +1 -1
img_nsfw.py CHANGED
@@ -87,4 +87,4 @@ def check_nsfw(img, pipe):
87
  from torch.cuda.amp import autocast
88
  with autocast():
89
  _, nsfw_tags = pipe.safety_checker.forward(clip_input=safety_checker_input.pixel_values, images=img)
90
- return nsfw_tags['bad_tags']
 
87
  from torch.cuda.amp import autocast
88
  with autocast():
89
  _, nsfw_tags = pipe.safety_checker.forward(clip_input=safety_checker_input.pixel_values, images=img)
90
+ return nsfw_tags[0]['bad_tags']