fix model device

app.py

@@ -84,11 +84,18 @@ def predict(model_name, pairs_file, sequence_file, progress = gr.Progress()):
     try:
         run_id = uuid4()
         device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+        use_cuda = torch.cuda.is_available()
 
         # gr.Info("Loading model...")
         _ = lm_embed("M", use_cuda = (device.type == "cuda"))
         
-        model = DSCRIPTModel.from_pretrained(model_map[model_name], use_cuda=torch.cuda.is_available())
+        model = DSCRIPTModel.from_pretrained(model_map[model_name], use_cuda=use_cuda)
+        if use_cuda:
+            model = model.to(device)
+            model.use_cuda = True
+        else:
+            model = model.to("cpu")
+            model.use_cuda = False
 
         # gr.Info("Loading files...")
         try: