Fix PSPNet ONNX export issue (#1270)

* fix(pspnet): fix ONNX export by replacing AdaptiveAvgPool2d with interpolate

* fix: restore state_dict compatibility and handle ONNX export in forward

* Update segmentation_models_pytorch/decoders/pspnet/decoder.py

* fix: support ONNX export and TorchScript for PSPNet

---------

Co-authored-by: Raphael Lapertot <raphael.lapertot@gmail.com>
Co-authored-by: Pavel Iakubovskii <qubvel@gmail.com>

Author: lapertor

segmentation_models_pytorch/decoders/pspnet/decoder.py

@@ -20,6 +20,8 @@ def __init__(
         if pool_size == 1:
             use_norm = "identity"  # PyTorch does not support BatchNorm for 1x1 shape
 
+        self.pool_size = pool_size
+
         self.pool = nn.Sequential(
             nn.AdaptiveAvgPool2d(output_size=(pool_size, pool_size)),
             modules.Conv2dReLU(
@@ -29,7 +31,19 @@ def __init__(
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         height, width = x.shape[2:]
-        x = self.pool(x)
+
+        if torch.jit.is_scripting():
+            # TorchScript path: use standard AdaptiveAvgPool2d via self.pool
+            x = self.pool(x)
+        elif torch.onnx.is_in_onnx_export():
+            # ONNX export path: AdaptiveAvgPool2d is often problematic during export.
+            # Using F.interpolate with 'area' mode provides the same mathematical result
+            # (average pooling) while being more robustly supported.
+            x = F.interpolate(x, size=(self.pool_size, self.pool_size), mode="area")
+            x = self.pool[1](x)  # use only ConvRelu block from pool
+        else:
+            x = self.pool(x)
+
         x = F.interpolate(x, size=(height, width), mode="bilinear", align_corners=True)
         return x