Update src/run.py

src/run.py

@@ -49,10 +49,10 @@ class NLP_classification():
         
     def training(self, epochs=50, batch_size=4, lr=1e-5, dropout=0.1, data_cut=None, early_stop_count=10, 
                 wandb_log=False, wandb_project=None, wandb_group=None, wandb_name=None, wandb_memo=None):
-        os.environ["CUDA_VISIBLE_DEVICES"]= "{0}".format(int(self.gpu_num))
-        #device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        device = torch.device('cuda:{0}'.format(int(self.gpu_num)))
-        torch.cuda.set_device(device)
+        #os.environ["CUDA_VISIBLE_DEVICES"]= "{0}".format(int(self.gpu_num))
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        #device = torch.device('cuda:{0}'.format(int(self.gpu_num)))
+        #torch.cuda.set_device(device)
         set_seed(self.random_state)
         torch.set_num_threads(10)
         
@@ -149,7 +149,7 @@ class NLP_classification():
         
         
     def prediction(self, selected_model=None, batch_size=8):
-        os.environ["CUDA_VISIBLE_DEVICES"]= "{0}".format(int(self.gpu_num))
+        #os.environ["CUDA_VISIBLE_DEVICES"]= "{0}".format(int(self.gpu_num))
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         set_seed(self.random_state)
         torch.set_num_threads(10)
@@ -192,10 +192,10 @@ class NLP_classification():
     
     
     def get_embedding(self, selected_model=None, batch_size=8, return_hidden=True, return_hidden_pretrained=False):
-        os.environ["CUDA_VISIBLE_DEVICES"]= "{0}".format(int(self.gpu_num))
-        #device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        device = torch.device('cuda:{0}'.format(int(self.gpu_num)))
-        torch.cuda.set_device(device)
+        #os.environ["CUDA_VISIBLE_DEVICES"]= "{0}".format(int(self.gpu_num))
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        #device = torch.device('cuda:{0}'.format(int(self.gpu_num)))
+        #torch.cuda.set_device(device)
         set_seed(self.random_state)
         torch.set_num_threads(10)
         task_type=self.task_type