Update app.py

app.py

@@ -21,9 +21,9 @@ model.to('cuda')
 # model_esm.to('cuda')
 # model_esm.eval()
 tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t36_3B_UR50D")
-model = EsmModel.from_pretrained("facebook/esm2_t36_3B_UR50D")
-model.to('cuda')
-model.eval()
+model_esm = EsmModel.from_pretrained("facebook/esm2_t36_3B_UR50D")
+model_esm.to('cuda')
+model_esm.eval()
 
 @spaces.GPU
 def generate_caption(protein, prompt):
@@ -94,9 +94,9 @@ def generate_caption(protein, prompt):
                     result["contacts"] = contacts[i, : truncate_len, : truncate_len].clone()
             esm_emb = result['representations'][36]
     '''
-    inputs = tokenizer([protein], return_tensors="pt", padding=True, truncation=True)
+    inputs = tokenizer([protein], return_tensors="pt", padding=True, truncation=True).to('cuda')
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = model_esm(**inputs)
     esm_emb = outputs.last_hidden_state.detach()[0]
     
     print("esm embedding generated")