Update app.py

app.py

@@ -37,15 +37,17 @@ def generate_caption(protein, prompt):
     print("dataset prepared")
     batches = dataset.get_batch_indices(toks_per_batch, extra_toks_per_seq=1)
     print("batches prepared")
+
+    model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
+    model_esm.to('cuda')
+    model_esm.eval()
+    
     data_loader = torch.utils.data.DataLoader(
         dataset, collate_fn=alphabet.get_batch_converter(truncation_seq_length), batch_sampler=batches
     )
     print(f"Read sequences")
     return_contacts = "contacts" in include
 
-    model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
-    model_esm.to('cuda')
-    model_esm.eval()
     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]