Update README.md

README.md

@@ -16,11 +16,37 @@ Deep learning models have shown remarkable performance in electrocardiogram (ECG
 
 ## Models
 This repository contains:
+- SMALL/BASE/LARGE HuBERTECG model sizes ready to be fine-tuned on any downstream dataset or to be used as feature extractor
 - SMALL/BASE/LARGE HuBERTECG model sizes fine-tuned on Cardio-Learning for a more disease-oriented baseline to futher fine-tune.
 
 Cardio-Learning is the name we gave to the union of several 12-lead ECG datasets including PTB, PTB-XL, CPSC, CPSC-Extra, Georgia, Chapman, Ningbo, SPH, CODE, SaMi-Trop, Hefei.
 This dataset, counting 2.4 million ECGs from millions of patients in 4 countries, encompasses 164 different heart-related conditions for which the ECG is either the primary or a supportive diagnostic tool, or is used to estimate the risk of future adverse cardiovascular events.
 
+## Usage
+```
+import torch
+from hubert_ecg import HuBERTECG
+
+path = "path/to/your/hubert-ecg-model.pt"
+checkpoint = torch.load(path, map_location='cpu')
+config = checkpoint['model_config']
+hubert_ecg = HuBERTECG(config)
+hubert_ecg.load_state_dict(checkpoint['model_state_dict']) # pre-trained model ready to be fine-tuned or used as feature extractor
+```
+
+```
+import torch
+from hubert_ecg import HuBERTECG
+from hubert_ecg_classification import HuBERTForECGClassification
+
+path = "path/to/your/finetuned-hubert-ecg-model.pt"
+checkpoint = torch.load(path, map_location='cpu')
+config = checkpoint['model_config']
+hubert_ecg = HuBERTECG(config)
+hubert_ecg = HuBERTForECGClassification(hubert)
+hubert_ecg.load_state_dict(checkpoint['model_state_dict']) # fine-tuned model ready to be used or further fine-tuned
+```
+
 ## 📚 Citation
 If you use our models or find our work useful, please consider citing us:
 ```