hasnanmr commited on
Commit
92e43b3
1 Parent(s): bb96f79

update readme

Browse files
Files changed (1) hide show
  1. README.md +18 -1
README.md CHANGED
@@ -2,4 +2,21 @@
2
  datasets:
3
  - garythung/trashnet
4
  pipeline_tag: image-classification
5
- ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
  datasets:
3
  - garythung/trashnet
4
  pipeline_tag: image-classification
5
+ ---
6
+
7
+ to load this state model use this step:
8
+ #define the model
9
+ model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
10
+ for param in model_resnet.parameters():
11
+ param.requires_grad = False
12
+
13
+ num_ftrs = model_resnet.fc.in_features
14
+ model_resnet.fc = nn.Linear(num_ftrs, 6)
15
+
16
+ # Load the weights
17
+ state_dict = torch.load('trashnet_resnet50.pth')
18
+ model.load_state_dict(state_dict)
19
+
20
+ # Switch to evaluation mode
21
+ model.eval()
22
+