Update README.md

README.md

@@ -34,10 +34,20 @@ This model is a small timm/vit_base_patch16_224.orig_in21k_ft_in1k trained on sv
 Use the code below to get started with the model.
 
 ```python
-import detectors
 import timm
-
-model = timm.create_model("vit_base_patch16_224_in21k_ft_svhn", pretrained=True)
+import torch
+from torch import nn
+
+model = timm.create_model("timm/vit_base_patch16_224.orig_in21k_ft_in1k",
+pretrained=False)
+model.head = nn.Linear(model.head.in_features, 10)
+model.load_state_dict(
+    torch.hub.load_state_dict_from_url(
+        "https://huggingface.co/edadaltocg/vit_base_patch16_224_in21k_ft_svhn/resolve/main/pytorch_model.bin",
+        map_location="cpu",
+        file_name="vit_base_patch16_224_in21k_ft_svhn.pth",
+    )
+)
 ```
 
 ## Training Data