edadaltocg commited on
Commit
e8c9d8d
1 Parent(s): c584068

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +13 -3
README.md CHANGED
@@ -34,10 +34,20 @@ This model is a small timm/vit_base_patch16_224.orig_in21k_ft_in1k trained on sv
34
  Use the code below to get started with the model.
35
 
36
  ```python
37
- import detectors
38
  import timm
39
-
40
- model = timm.create_model("vit_base_patch16_224_in21k_ft_svhn", pretrained=True)
 
 
 
 
 
 
 
 
 
 
 
41
  ```
42
 
43
  ## Training Data
 
34
  Use the code below to get started with the model.
35
 
36
  ```python
 
37
  import timm
38
+ import torch
39
+ from torch import nn
40
+
41
+ model = timm.create_model("timm/vit_base_patch16_224.orig_in21k_ft_in1k",
42
+ pretrained=False)
43
+ model.head = nn.Linear(model.head.in_features, 10)
44
+ model.load_state_dict(
45
+ torch.hub.load_state_dict_from_url(
46
+ "https://huggingface.co/edadaltocg/vit_base_patch16_224_in21k_ft_svhn/resolve/main/pytorch_model.bin",
47
+ map_location="cpu",
48
+ file_name="vit_base_patch16_224_in21k_ft_svhn.pth",
49
+ )
50
+ )
51
  ```
52
 
53
  ## Training Data