fix dygraph re5_head bug, test=dygraph (open-mmlab#2088)

liuhuiCNN · Jan 18, 2021 · 2bdb9a8 · 2bdb9a8
1 parent e7dbeb0
commit 2bdb9a8
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/dygraph/ppdet/modeling/backbones/resnet.py b/dygraph/ppdet/modeling/backbones/resnet.py
@@ -530,14 +530,15 @@ def forward(self, inputs):
 
 @register
 class Res5Head(nn.Layer):
-    def __init__(self, feat_in=1024, feat_out=512):
+    def __init__(self, depth=50, feat_in=1024, feat_out=512):
         super(Res5Head, self).__init__()
         na = NameAdapter(self)
         self.res5_conv = []
         self.res5 = self.add_sublayer(
             'res5_roi_feat',
             Blocks(
-                feat_in, feat_out, count=3, name_adapter=na, stage_num=5))
+                depth, feat_in, feat_out, count=3, name_adapter=na,
+                stage_num=5))
         self.feat_out = feat_out * 4
 
     def forward(self, roi_feat, stage=0):