diff --git a/Experiments/evaluate_tabular.py b/Experiments/evaluate_tabular.py index ba0654f071cdb96cfa72145fe048da4b088bc16b..1cb0fec9875b2e0b7fad3fd42acb2bb4f9057648 100644 --- a/Experiments/evaluate_tabular.py +++ b/Experiments/evaluate_tabular.py @@ -39,9 +39,25 @@ train_data, test_data = load_data() input_dim = train_data[0][0].numel() output_dim = train_data[0][1].numel() +try: + gpu_number = eiv_conf_dict["gpu_number"] + device = torch.device(f'cuda:{gpu_number}') + try: + torch.tensor([0.0]).to(device) + except RuntimeError: + if torch.cuda.is_available(): + print('Switched to GPU 0') + device = torch.device('cuda:0') + else: + print('No cuda available, using CPU') + device = torch.device('cpu') +except KeyError: + device = torch.device('cpu') + + def collect_metrics(x,y, seed=0, noneiv_number_of_draws=100, eiv_number_of_draws=[100,5], - decouple_dimensions=False, device=torch.device('cuda:1'), + decouple_dimensions=False, device=device, scale_outputs=scale_outputs): """ Compute various metrics for EiV and non-EiV. Will be returned as diff --git a/Experiments/train_eiv.py b/Experiments/train_eiv.py index eb8d9d88292dd63130b38bc78db4eb1ed1db9014..68229469c167fc934a06bbf61e7c5673a3005b12 100644 --- a/Experiments/train_eiv.py +++ b/Experiments/train_eiv.py @@ -55,8 +55,16 @@ print(f"Training on {long_dataname} data") try: gpu_number = conf_dict["gpu_number"] - device = torch.device(f'cuda:{gpu_number}' if torch.cuda.is_available() - else 'cpu') + device = torch.device(f'cuda:{gpu_number}') + try: + torch.tensor([0.0]).to(device) + except RuntimeError: + if torch.cuda.is_available(): + print('Switched to GPU 0') + device = torch.device('cuda:0') + else: + print('No cuda available, using CPU') + device = torch.device('cpu') except KeyError: device = torch.device('cpu') diff --git a/Experiments/train_noneiv.py b/Experiments/train_noneiv.py index 316dda1e8f5fbf42f46ae9135dadd5c803623df8..ee5687b268d762aeadac294ebdd6e43c37a25149 100644 --- a/Experiments/train_noneiv.py +++ b/Experiments/train_noneiv.py @@ -54,8 +54,16 @@ print(f"Training on {long_dataname} data") try: gpu_number = conf_dict["gpu_number"] - device = torch.device(f'cuda:{gpu_number}' if torch.cuda.is_available() - else 'cpu') + device = torch.device(f'cuda:{gpu_number}') + try: + torch.tensor([0.0]).to(device) + except RuntimeError: + if torch.cuda.is_available(): + print('Switched to GPU 0') + device = torch.device('cuda:0') + else: + print('No cuda available, using CPU') + device = torch.device('cpu') except KeyError: device = torch.device('cpu')