Update README.md

README.md

@@ -35,16 +35,17 @@ is not validated, it was primarily used a pre-training task.
 
 ### How to use
 
-You can use this model directly by loading via the Structure2EnergyandForces task:
+Here is how to use the model to extract features from the pre-trained backbone:
 
 ```python
->>> from transformers import AutoModel
-```
+import torch
+from transformers import AutoModel
+model = AutoModel.from_pretrained("vector-institute/atomformer-base", trust_remote_code=True)
 
-Here is how to use this model to get the features of a given atomistic graph in PyTorch:
+input_ids, coords, attn_mask = torch.randint(0, 100, (1, 10)), torch.randn(1, 10, 3), torch.ones(1, 10)
 
-```python
-from transformers import AutoModel
+output = model(input_ids, coords=coords, attention_mask=attention_mask)
+output[0].shape #  (torch.Size([1, 10, 768])
 ```
 
 
@@ -54,9 +55,6 @@ AtomFormer is trained on an aggregated S2EF dataset from multiple sources such a
 with structures and energies/forces for pre-training.  The pre-training data includes total energies and formation
 energies but trains using formation energy (which isn't included for OC22, indicated by "has_formation_energy" column).
 
-## Training procedure
-
-
 
 ### Preprocessing