diff --git a/api/onnx_web/chain/edit_safety.py b/api/onnx_web/chain/edit_safety.py
index afcb7fb6..d4635b2e 100644
--- a/api/onnx_web/chain/edit_safety.py
+++ b/api/onnx_web/chain/edit_safety.py
@@ -39,7 +39,7 @@ class EditSafetyStage(BaseStage):
             deep_danbooru_model,
         )
 
-        server.cache.set(ModelTypes.safety, cache_key)
+        server.cache.set(ModelTypes.safety, cache_key, nsfw_checker)
 
         return nsfw_checker