提交 9e5c26da 编写于 作者: Y Yang Zhang 提交者: qingqing01

Fix ResNet assertion error when `feature_maps` is an integer (#2508)

上级 23441fa3
......@@ -56,6 +56,9 @@ class ResNet(object):
feature_maps=[2, 3, 4, 5]):
super(ResNet, self).__init__()
if isinstance(feature_maps, Integral):
feature_maps = [feature_maps]
assert depth in [18, 34, 50, 101, 152], \
"depth {} not in [18, 34, 50, 101, 152]"
assert variant in ['a', 'b', 'c', 'd'], "invalid ResNet variant"
......@@ -70,8 +73,6 @@ class ResNet(object):
self.freeze_norm = freeze_norm
self.variant = variant
self._model_type = 'ResNet'
if isinstance(feature_maps, Integral):
feature_maps = [feature_maps]
self.feature_maps = feature_maps
self.depth_cfg = {
18: ([2, 2, 2, 2], self.basicblock),
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册