Merge 8d027df229 into ee7f2e8850

mmpretrain/models/necks/simmim_neck.py

@@ -15,14 +15,18 @@ class SimMIMLinearDecoder(BaseModule):
     Args:
         in_channels (int): Channel dimension of the feature map.
         encoder_stride (int): The total stride of the encoder.
+        target_channels (int): Channel dimensions of original image.
     """
 
-    def __init__(self, in_channels: int, encoder_stride: int) -> None:
+    def __init__(self,
+                 in_channels: int,
+                 encoder_stride: int,
+                 target_channels: int = 3) -> None:
         super().__init__()
         self.decoder = nn.Sequential(
             nn.Conv2d(
                 in_channels=in_channels,
-                out_channels=encoder_stride**2 * 3,
+                out_channels=encoder_stride**2 * target_channels,
                 kernel_size=1),
             nn.PixelShuffle(encoder_stride),
         )