update code

README.md

@@ -26,13 +26,13 @@ in terms of top-1 accuracy, efficiency, and robustness on my dataset and [CMAD b
 
 <br>
 
-### Load 1000 Class PyTorch Jit Model
+### Load PyTorch Jit Model with 1000 Classes
 ```python
 from transformers import AutoModel
 model = AutoModel.from_pretrained("FredZhang7/efficientnetv2.5_rw_s", trust_remote_code=True)
 ```
 
-### Prepare Model for Training
+### Load Model with Custom Classes
 To change the number of classes, replace the linear classification layer. 
 Here's an example of how to convert the architecture into a trainable model.
 ```bash