Save base_channels for UNet (#450)

Signed-off-by: lizz <lizz@sensetime.com>
2021-03-31 00:52:08 +08:00 · 2021-03-31 00:52:08 +08:00 · 455d163217
commit 455d163217
parent 33e8357c62
1 changed files with 1 additions and 0 deletions
--- a/mmseg/models/backbones/unet.py
+++ b/mmseg/models/backbones/unet.py
@ -329,6 +329,7 @@ class UNet(nn.Module):
        self.strides = strides
        self.downsamples = downsamples
        self.norm_eval = norm_eval
+        self.base_channels = base_channels

        self.encoder = nn.ModuleList()
        self.decoder = nn.ModuleList()