def predict(self, x):
        """Inference method. Switch model to `eval` mode, call `.forward(x)` with `torch.inference_mode()`

        Args:
            x: 4D torch tensor with shape (batch_size, channels, height, width)

        Return:
            prediction: 4D torch tensor with shape (batch_size, classes, height, width)

        """
        if self.training:
            self.eval()
        x = self(x)
        return x
Inference method. Switch model to `eval` mode, call `.forward(x)` with `torch.inference_mode()`

        Args:
            x: 4D torch tensor with shape (batch_size, channels, height, width)

        Return:
            prediction: 4D torch tensor with shape (batch_size, classes, height, width)

        
segmentation_models_pytorch/base/model.py