Skip to content

Commit fd37e1b

Browse files
leaderj1001ruotianluo
authored andcommitted
[Update] MobileNet-v1 load_pretrained_cnn function
When I tried to train the MobileNet-v1 pretrained model, I faced an error. So, I fixed it.
1 parent 1587ef1 commit fd37e1b

File tree

1 file changed

+1
-2
lines changed

1 file changed

+1
-2
lines changed

lib/nets/mobilenet_v1.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -289,6 +289,5 @@ def set_bn_eval(m):
289289
def load_pretrained_cnn(self, state_dict):
290290
print('Warning: No available pretrained model yet')
291291
self.mobilenet.load_state_dict({
292-
k: state_dict['features.' + k]
293-
for k in list(self.mobilenet.state_dict())
292+
k: state_dict['features.' + k] for k in list(self.mobilenet.state_dict()) if 'features.' + k in state_dict
294293
})

0 commit comments

Comments
 (0)