Make norm layer as parameters in models.detection.backbone_utils.py (#…

…2081)
pytorch · Apr 9, 2020 · c89156b · c89156b
1 parent 684f48d
commit c89156b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchvision/models/detection/backbone_utils.py b/torchvision/models/detection/backbone_utils.py
@@ -41,10 +41,10 @@ def forward(self, x):
         return x
 
 
-def resnet_fpn_backbone(backbone_name, pretrained):
+def resnet_fpn_backbone(backbone_name, pretrained, norm_layer=misc_nn_ops.FrozenBatchNorm2d):
     backbone = resnet.__dict__[backbone_name](
         pretrained=pretrained,
-        norm_layer=misc_nn_ops.FrozenBatchNorm2d)
+        norm_layer=norm_layer)
     # freeze layers
     for name, parameter in backbone.named_parameters():
         if 'layer2' not in name and 'layer3' not in name and 'layer4' not in name: