Update README.md

README.md

@@ -27,9 +27,13 @@ ESM_DIM = 1280
 SAE_DIM = 4096
 LAYER = 24
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # Load ESM model
 tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
 esm_model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D")
+esm_model.to(device)
+esm_model.eval()
 
 # Load SAE model
 checkpoint_path = hf_hub_download(
@@ -38,11 +42,13 @@ checkpoint_path = hf_hub_download(
 )
 sae_model = SparseAutoencoder(ESM_DIM, SAE_DIM)
 sae_model.load_state_dict(load_file(checkpoint_path))
+sae_model.to(device)
+sae_model.eval()
 ```
 
 ESM -> SAE inference on an amino acid sequence of length `L`
 ```
-seq = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVVAAIVQDIAYLRSLGYNIVATPRGYVLAGG"
+seq = "TTCCPSIVARSNFNVCRLPGTPEALCATYTGCIIIPGATCPGDYAN"
 
 # Tokenize sequence and run ESM inference
 inputs = tokenizer(seq, padding=True, return_tensors="pt")