Update README.md

README.md

@@ -1,12 +1,14 @@
 ---
+license: apache-2.0
+pipeline_tag: image-feature-extraction
 tags:
 - model_hub_mixin
 - pytorch_model_hub_mixin
-license: apache-2.0
 ---
+
 # Communication-Inspired Tokenization for Structured Image Representations
-</h1>
-  <p align="left">
+
+<p align="left">
     <a href="https://araachie.github.io">Aram Davtyan</a> •
     <a href="https://www.cvg.unibe.ch/people/sahin">Yusuf Sahin</a> •
     <a href="https://people.epfl.ch/yasaman.haghighi?lang=en">Yasaman Haghighi</a> •
@@ -14,23 +16,25 @@ license: apache-2.0
     <a href="https://www.cvg.unibe.ch/people/acuaviva">Pablo Acuaviva</a> •
     <a href="https://people.epfl.ch/alexandre.alahi?lang=en">Alexandre Alahi</a> •
     <a href="https://www.cvg.unibe.ch/people/favaro">Paolo Favaro</a>
-  </p>
+</p>
+
+Official pre-trained models for the paper: [Communication-Inspired Tokenization for Structured Image Representations](https://arxiv.org/abs/2602.20731).
 
-Official pre-trained models for the paper: https://arxiv.org/abs/2602.20731
+[[Website](https://araachie.github.io/comit/)] [[Code](https://github.com/Araachie/comit)] [[Paper](https://arxiv.org/abs/2602.20731)]
 
-Project's website: https://araachie.github.io/comit/
-  
 ## Installation
 
-Follow the instructions at https://github.com/Araachie/comit
+Follow the instructions at [https://github.com/Araachie/comit](https://github.com/Araachie/comit)
 
 ## Usage
 
 Example usage, downloading `COMiT-B` from the Hugging Face Hub:
 
 ```python
+import torch
 from comit import COMiT
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
 model = COMiT.from_pretrained('cvg-unibe/comit-b')
 model.eval().to(device)
 ```
@@ -84,7 +88,7 @@ with torch.no_grad():
 ## Licensing
 
 Unless otherwise noted, the model weights are licensed under Apache license 2.0.
-For the code licensing, see https://github.com/Araachie/comit?tab=readme-ov-file#licensing
+For the code licensing, see [https://github.com/Araachie/comit?tab=readme-ov-file#licensing](https://github.com/Araachie/comit?tab=readme-ov-file#licensing)
 
 ## Citation