[Support PyTorch 1.8] Remove inference mode (#707)

2022-10-03 22:14:58 +02:00 · 2022-10-03 22:14:58 +02:00 · b35bac4d3b
parent 688031c592
commit b35bac4d3b
1 changed files with 1 additions and 1 deletions
--- a/src/diffusers/pipelines/stable_diffusion/safety_checker.py
+++ b/src/diffusers/pipelines/stable_diffusion/safety_checker.py
@ -79,7 +79,7 @@ class StableDiffusionSafetyChecker(PreTrainedModel):

        return images, has_nsfw_concepts

-    @torch.inference_mode()
+    @torch.no_grad()
    def forward_onnx(self, clip_input: torch.FloatTensor, images: torch.FloatTensor):
        pooled_output = self.vision_model(clip_input)[1]  # pooled_output
        image_embeds = self.visual_projection(pooled_output)