diff --git a/Experiments/plot_real_diagonal.py b/Experiments/plot_diagonal_uncertainties.py
similarity index 89%
rename from Experiments/plot_real_diagonal.py
rename to Experiments/plot_diagonal_uncertainties.py
index d2abda6db2c527edd0731c0a544d1771b425f271..5401a44a1e7665d65e2176ecd1f196ffd52a1c05 100644
--- a/Experiments/plot_real_diagonal.py
+++ b/Experiments/plot_diagonal_uncertainties.py
@@ -62,19 +62,7 @@ train_data, test_data = load_data(normalize=normalize)
 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:
+# do computations on cpu
 device = torch.device('cpu')
 
 
@@ -164,19 +152,19 @@ seed_list = range(noneiv_conf_dict["seed_range"][0],
 noneiv_uncertainties = 0
 eiv_uncertainties = 0
 number_of_seeds = len(seed_list)
-out_dim = 0
 number_of_steps = 100
 for seed in tqdm(seed_list):
     x_diagonal = create_diagonal(train=train_data, number_of_steps=number_of_steps)
-    results = collect_predictions(x_diagonal,
-            seed=seed)
-    noneiv_uncertainties += 1/number_of_seeds * results['noneiv']['uncertainties'][..., out_dim]
-    eiv_uncertainties += 1/number_of_seeds * results['eiv']['uncertainties'][..., out_dim]
+    results = collect_predictions(x_diagonal, seed=seed)
+    noneiv_uncertainties += 1/number_of_seeds * results['noneiv']['uncertainties'].mean(dim=-1)
+    eiv_uncertainties += 1/number_of_seeds * results['eiv']['uncertainties'].mean(dim=-1)
 
 
 plt.figure(1)
 plt.clf()
-plot_x = torch.linspace(0,1, steps=number_of_steps)
+plot_x = torch.linspace(0, 1, steps=number_of_steps)
 plt.fill_between(plot_x, noneiv_uncertainties, color='b', alpha=0.5)
 plt.fill_between(plot_x,  eiv_uncertainties, color='r', alpha=0.5)
-plt.savefig('results/figures/intersection')
+plt.xlabel(r'$\lambda$')
+plt.ylabel(r'$u$')
+plt.savefig(f'results/figures/diagonal_uncertainties_{data}.pdf')