mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Fix model create fn not passing num_classes through. Fix #135

Browse Source
This commit is contained in:
Ross Wightman 2020-05-01 18:37:03 -07:00
parent 779cb0fcc0
commit ea300709f0
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
timm/models/pnasnet.py
View File

@ -390,7 +390,7 @@ def pnasnet5large(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
<https://arxiv.org/abs/1712.00559>`_ paper.
"""
default_cfg = default_cfgs['pnasnet5large']
model = PNASNet5Large(num_classes=1000, in_chans=in_chans, **kwargs)
model = PNASNet5Large(num_classes=num_classes, in_chans=in_chans, **kwargs)
model.default_cfg = default_cfg
if pretrained:
load_pretrained(model, default_cfg, num_classes, in_chans)

2
timm/models/res2net.py
View File

@ -229,7 +229,7 @@ def res2next50(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
default_cfg = default_cfgs['res2next50']
res2net_block_args = dict(scale=4)
model = ResNet(Bottle2neck, [3, 4, 6, 3], base_width=4, cardinality=8,
num_classes=1000, in_chans=in_chans, block_args=res2net_block_args, **kwargs)
num_classes=num_classes, in_chans=in_chans, block_args=res2net_block_args, **kwargs)
model.default_cfg = default_cfg
if pretrained:
load_pretrained(model, default_cfg, num_classes, in_chans)