Update README.md

README.md

@@ -4,6 +4,7 @@ base_model:
 - google-bert/bert-base-uncased
 ---
 
+
 You can load the laff model using the following steps: 
 
 ```
@@ -11,6 +12,11 @@ import torch
 from transformers import BertTokenizer, BertForSequenceClassification
 from base_models import  BinaryClassificationBertModel
 
+from huggingface_hub import hf_hub_download
+
+file_path = hf_hub_download(repo_id="mandeep-rathee/laff-model", filename="bert-base-laff.pth")
+
+
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 base_model_name = "bert-base-uncased" 
@@ -18,6 +24,6 @@ tokenizer = BertTokenizer.from_pretrained(base_model_name, torch_dtype=torch.flo
 base_model = BertForSequenceClassification.from_pretrained(base_model_name, num_labels=1,torch_dtype=torch.float16)
 
 model = BinaryClassificationBertModel(base_model)
-model.load_state_dict(torch.load("laff_model/bert-base-laff.pth"))
+model.load_state_dict(torch.load(file_path, map_location=device))
 model.to(device)
 ```