Update README.md

README.md

@@ -27,7 +27,12 @@ gpu_id = 0
 device = get_torch_device(gpu_id)
 
 # load the model
-model = HierarchyTransformer.load_pretrained('Hierarchy-Transformers/HiT-MiniLM-L12-WordNet', device)
+revision = "main"  # change for a different version
+model = HierarchyTransformer.from_pretrained(
+  model_name_or_path='Hierarchy-Transformers/HiT-MiniLM-L12-WordNetNoun',
+  revision=revision
+  device=device
+)
 
 # entity names to be encoded.
 entity_names = ["computer", "personal computer", "fruit", "berry"]