support loading deit weights (#538)

open-mmlab · May 9, 2021 · 584f5a7 · 584f5a7
1 parent db44d16
commit 584f5a7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/mmseg/models/backbones/vit.py b/mmseg/models/backbones/vit.py
@@ -325,6 +325,8 @@ def init_weights(self, pretrained=None):
             checkpoint = _load_checkpoint(pretrained, logger=logger)
             if 'state_dict' in checkpoint:
                 state_dict = checkpoint['state_dict']
+            elif 'model' in checkpoint:
+                state_dict = checkpoint['model']
             else:
                 state_dict = checkpoint
 

diff --git a/mmseg/models/necks/multilevel_neck.py b/mmseg/models/necks/multilevel_neck.py
@@ -54,7 +54,6 @@ def __init__(self,
 
     def forward(self, inputs):
         assert len(inputs) == len(self.in_channels)
-        print(inputs[0].shape)
         inputs = [
             lateral_conv(inputs[i])
             for i, lateral_conv in enumerate(self.lateral_convs)