fix bug

haojiepan1 · Nov 22, 2021 · 5e32826 · 5e32826
1 parent 4319ee7
commit 5e32826
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/bart_auto_title_train.py b/examples/bart_auto_title_train.py
@@ -104,10 +104,10 @@ def padding(indice, max_length, pad_idx=0):
 class Trainer:
     def __init__(self):
         # 加载数据
-        self.sents_src, self.sents_tgt = read_file(src_dir, tgt_dir)
+        self.sents_src, self.sents_tgt = read_file()
 
         # 判断是否有可用GPU
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
         print("device: " + str(self.device))
         # 定义模型
         self.model = ExtendModel(model, tokenizer=tokenizer, bos_id=word2idx["[CLS]"], eos_id=word2idx["[SEP]"], device=self.device)