Skip to content

Commit a520da9

Browse files
committed
Update tresnet features_info for v2
1 parent c8ab747 commit a520da9

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

timm/models/tresnet.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -200,8 +200,8 @@ def __init__(
200200

201201
self.feature_info = [
202202
dict(num_chs=self.planes, reduction=2, module=''), # Not with S2D?
203-
dict(num_chs=self.planes, reduction=4, module='body.layer1'),
204-
dict(num_chs=self.planes * 2, reduction=8, module='body.layer2'),
203+
dict(num_chs=self.planes * (Bottleneck.expansion if v2 else 1), reduction=4, module='body.layer1'),
204+
dict(num_chs=self.planes * 2 * (Bottleneck.expansion if v2 else 1), reduction=8, module='body.layer2'),
205205
dict(num_chs=self.planes * 4 * Bottleneck.expansion, reduction=16, module='body.layer3'),
206206
dict(num_chs=self.planes * 8 * Bottleneck.expansion, reduction=32, module='body.layer4'),
207207
]

0 commit comments

Comments
 (0)