Skip to content
Snippets Groups Projects
Commit 61e39589 authored by Schneider Leo's avatar Schneider Leo
Browse files

fix : resnet50 linear layer shape

parent ea13212d
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Showing
with 1 addition and 1 deletion
models/model.py
+ 1
− 1
View file @ 61e39589
...@@ -324,7 +324,7 @@ class Classification_model_duo_pretrained(nn.Module): ...@@ -324,7 +324,7 @@ class Classification_model_duo_pretrained(nn.Module):
if model =='ResNet50': if model =='ResNet50':
self.im_encoder = torch.hub.load('pretrained_model/pytorch_vision_v0.10.0/', 'resnet50', pretrained=True,source='local') self.im_encoder = torch.hub.load('pretrained_model/pytorch_vision_v0.10.0/', 'resnet50', pretrained=True,source='local')
self.im_encoder.fc = torch.nn.Linear(2000, self.n_class) self.im_encoder.fc = torch.nn.Linear(2048, self.n_class)
#freeze backbone #freeze backbone
for n, p in self.im_encoder.named_parameters(): for n, p in self.im_encoder.named_parameters():
if n not in ['fc.weight', 'fc.bias']: if n not in ['fc.weight', 'fc.bias']:
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment