self.device set to gpu

new_test_saved_finetuned_model.py

@@ -47,8 +47,8 @@ class BERTFineTuneTrainer:
 
         # Setup cuda device for BERT training, argument -c, --cuda should be true
         # cuda_condition = torch.cuda.is_available() and with_cuda
-        # self.device = torch.device("cuda:0" if cuda_condition else "cpu")
-        self.device = torch.device("cpu") #torch.device("cuda:0" if cuda_condition else "cpu")
+        self.device = torch.device("cuda:0" if cuda_condition else "cpu")
+        #self.device = torch.device("cpu") #torch.device("cuda:0" if cuda_condition else "cpu")
         # print(cuda_condition, " Device used = ", self.device)
         print(" Device used = ", self.device)