Update app.py

app.py

@@ -10,6 +10,7 @@ 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()
@@ -18,7 +19,7 @@ model_esm = model_esm.cuda()
 # Load the model
 model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
 model.load_checkpoint("model/checkpoint_mf2.pth")
-model.to('cuda')
+model = model.cuda()
 
 
 @spaces.GPU