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

dereyly · Oct 3, 2022 · b35bac4 · b35bac4
1 parent 688031c
commit b35bac4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/diffusers/pipelines/stable_diffusion/safety_checker.py b/src/diffusers/pipelines/stable_diffusion/safety_checker.py
@@ -79,7 +79,7 @@ def forward(self, clip_input, images):
 
         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)