Handling of unaccessible GPU added

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

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')
 

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')