fix device issue: cpu

src/check_client/fact_checker.py

@@ -56,9 +56,9 @@ class FactChecker:
         self.tokenizer = None
         self.model = None
         self.args = args
+        self.args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.ckpt = args.fc_dir if fc_ckpt_dir is None else fc_ckpt_dir
         self.mask_rate = mask_rate
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
         logger.info('Initializing fact checker.')
         self._prepare_ckpt(self.args.model_name_or_path, self.ckpt)
@@ -98,7 +98,7 @@ class FactChecker:
         dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)
 
         with torch.no_grad():
-            self.model.to(self.device)
+            self.model.to(self.args.device)
             self.model.eval()
             iter = tqdm(dataloader, desc="Fact Checking") if verbose else dataloader
             _, y_predicted, z_predicted, m_attn, mask = \

src/check_client/train.py

@@ -267,7 +267,7 @@ def do_evaluate(dataloader, model, args, during_training=False, with_label=True)
     mask = []
     for i, batch in enumerate(dataloader):
         model.eval()
-        batch = tuple(t.to(model.device) for t in batch)
+        batch = tuple(t.to(args.device) for t in batch)
         with torch.no_grad():
             inputs = {
                 "claim_input_ids": batch[0],