fixed bug

README.md

@@ -25,12 +25,15 @@ You can try out the MetaLATTE model directly in your browser:
 
 ```python
 from transformers import AutoTokenizer, AutoModel, AutoConfig
+from metalatte import MetaLATTEConfig, MetaLATTEModel
+AutoConfig.register("metalatte", MetaLATTEConfig)
+AutoModel.register(MetaLATTEConfig, MetaLATTEModel)
 
 tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
 config = AutoConfig.from_pretrained("ChatterjeeLab/MetaLATTE")
 model = AutoModel.from_pretrained("ChatterjeeLab/MetaLATTE", config=config)
 
-# Use the model
+model.eval()
 sequence = "AVYNIGWSFNVNGARGKSFRAGDVLVFKYIKGQHNVVAVNGRGYASCSAPRGARTYSSGQDRIKLTRGQNYFICSFPGHCGGGMKIAINAK"
 inputs = tokenizer(sequence, return_tensors="pt")
 raw_probs, predictions = model(**inputs)

config.json

@@ -2,7 +2,7 @@
     "architectures": [
       "MultitaskProteinModel"
     ],
-    "model_type": "esm",
+    "model_type": "metalatte",
     "num_labels": 15,
     "hidden_size": 1280,
     "num_hidden_layers": 33,