Update app.py

app.py

@@ -27,7 +27,7 @@ def convert_outputs_to_pdb(outputs):
     return pdbs
 
 def fold_prot_locally(sequence):
-    tokenized_input = tokenizer([sequence], return_tensors="pt", add_special_tokens=False)['input_ids']
+    tokenized_input = tokenizer([sequence], return_tensors="pt", add_special_tokens=False)['input_ids'].cuda()
 
     with torch.no_grad():
         output = model(tokenized_input)
@@ -35,7 +35,7 @@ def fold_prot_locally(sequence):
     return pdb
 
 tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
-model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True)
+model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True).cuda().half()
 
 iface = gr.Interface(fn=fold_prot_locally, inputs="text", outputs="text")
 iface.launch()