fix device issue: cpu

src/check_client/fact_checker.py

@@ -58,6 +58,7 @@ class FactChecker:
         self.args = args
         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)
@@ -97,7 +98,7 @@ class FactChecker:
         dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)
 
         with torch.no_grad():
-            self.model.to(self.args.device)
+            self.model.to(self.device)
             self.model.eval()
             iter = tqdm(dataloader, desc="Fact Checking") if verbose else dataloader
             _, y_predicted, z_predicted, m_attn, mask = \
@@ -192,8 +193,6 @@ if __name__ == '__main__':
 
     set_seed(args)
 
-    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
     if args.output == 'none':
         args.ckpt = args.ckpt[:-1] if args.ckpt.endswith('/') else args.ckpt
         base_name = os.path.basename(args.ckpt)

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(args.device) for t in batch)
+        batch = tuple(t.to(model.device) for t in batch)
         with torch.no_grad():
             inputs = {
                 "claim_input_ids": batch[0],