Fixing image segmentation with inference mode.

src/transformers/pipelines/base.py

@@ -862,17 +862,19 @@ def postprocess(self, model_outputs: ModelOutput, **postprocess_parameters: Dict
         """
         raise NotImplementedError("postprocess not implemented")
 
+    def get_inference_context(self):
+        inference_context = (
+            torch.inference_mode if version.parse(torch.__version__) >= version.parse("1.9.0") else torch.no_grad
+        )
+        return inference_context
+
     def forward(self, model_inputs, **forward_params):
         with self.device_placement():
             if self.framework == "tf":
                 model_inputs["training"] = False
                 model_outputs = self._forward(model_inputs, **forward_params)
             elif self.framework == "pt":
-                inference_context = (
-                    torch.inference_mode
-                    if version.parse(torch.__version__) >= version.parse("1.9.0")
-                    else torch.no_grad
-                )
+                inference_context = self.get_inference_context()
                 with inference_context():
                     model_inputs = self._ensure_tensor_on_device(model_inputs, device=self.device)
                     model_outputs = self._forward(model_inputs, **forward_params)

src/transformers/pipelines/image_segmentation.py

@@ -114,6 +114,9 @@ def __call__(self, *args, **kwargs) -> Union[Predictions, List[Prediction]]:
 
         return super().__call__(*args, **kwargs)
 
+    def get_inference_context(self):
+        return torch.no_grad
+
     def preprocess(self, image):
         image = self.load_image(image)
         target_size = torch.IntTensor([[image.height, image.width]])