diff --git a/modules/GSOC/E2_ESRGAN/lib/dataset.py b/modules/GSOC/E2_ESRGAN/lib/dataset.py
index be07dd414f0ad2bfb20443ed935e35b308226c16..29335c9629f947fde759e68459828cf574f6f812 100644
--- a/modules/GSOC/E2_ESRGAN/lib/dataset.py
+++ b/modules/GSOC/E2_ESRGAN/lib/dataset.py
@@ -136,6 +136,7 @@ f.close()
 mean_std_dct.update(mean_std_dct_l1b)
 mean_std_dct.update(mean_std_dct_l2)
 
+
 class OpdNpyDataset:
 
     def __init__(self, filenames, hr_size, lr_size, batch_size=128):
@@ -179,12 +180,12 @@ class OpdNpyDataset:
             fname = self.filenames[fi]
             data = np.load(fname)
 
-            refl = data[0, ]
+            refl = data[0, 0, :, :]
             refl = scale(refl, 'refl_0_65um_nom', mean_std_dct)
             refl = refl.astype(np.float32)
             refl_s.append(refl)
 
-            opd = data[1, ]
+            opd = data[0, 1, :, :]
             opd = scale(opd, 'cld_opd_dcomp', mean_std_dct)
             opd = opd.astype(np.float32)
             opd_s.append(opd)