Commit
•
ac86014
1
Parent(s):
be558de
Update README.md
Browse files
README.md
CHANGED
|
@@ -27,10 +27,7 @@ def fix_names(state_dict):
|
|
| 27 |
state_dict = {key.replace('module.', ''): value for (key, value) in state_dict.items()}
|
| 28 |
return state_dict
|
| 29 |
|
| 30 |
-
model = models.mobilenet_v2()
|
| 31 |
-
|
| 32 |
-
num_ftrs = model.classifier[1].in_features
|
| 33 |
-
model.classifier[1] = nn.Linear(num_ftrs, FOOD101_CLASSES)
|
| 34 |
|
| 35 |
if len(sys.argv) > 1:
|
| 36 |
checkpoint_path = sys.argv[1]
|
|
|
|
| 27 |
state_dict = {key.replace('module.', ''): value for (key, value) in state_dict.items()}
|
| 28 |
return state_dict
|
| 29 |
|
| 30 |
+
model = models.mobilenet_v2(num_classes=FOOD101_CLASSES)
|
|
|
|
|
|
|
|
|
|
| 31 |
|
| 32 |
if len(sys.argv) > 1:
|
| 33 |
checkpoint_path = sys.argv[1]
|