Update app.py

app.py

@@ -10,16 +10,14 @@ import gradio as gr
 from esm_scripts.extract import run_demo
 from esm import pretrained
 
-print(torch.cuda.get_device_name(0))
 
 model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
-model_esm.eval()
-model_esm = model_esm.cuda()
+model_esm.to('cuda')
 
 # Load the model
 model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
 model.load_checkpoint("model/checkpoint_mf2.pth")
-model = model.cuda()
+model.to('cuda')
 
 
 @spaces.GPU
@@ -50,6 +48,7 @@ def generate_caption(protein, prompt):
     return_contacts = "contacts" in include
     assert all(-(model_esm.num_layers + 1) <= i <= model_esm.num_layers for i in repr_layers)
     repr_layers = [(i + model_esm.num_layers + 1) % (model_esm.num_layers + 1) for i in repr_layers]
+    model_esm.eval()
     with torch.no_grad():
         for batch_idx, (labels, strs, toks) in enumerate(data_loader):
             print(