diff --git a/modules/GSOC/E2_ESRGAN/lib/dataset.py b/modules/GSOC/E2_ESRGAN/lib/dataset.py
index 29335c9629f947fde759e68459828cf574f6f812..034e3bdf60b72637e3528353e995a0148d7d36ec 100644
--- a/modules/GSOC/E2_ESRGAN/lib/dataset.py
+++ b/modules/GSOC/E2_ESRGAN/lib/dataset.py
@@ -180,12 +180,12 @@ class OpdNpyDataset:
             fname = self.filenames[fi]
             data = np.load(fname)
 
-            refl = data[0, 0, :, :]
+            refl = data[:, 0, :, :]
             refl = scale(refl, 'refl_0_65um_nom', mean_std_dct)
             refl = refl.astype(np.float32)
             refl_s.append(refl)
 
-            opd = data[0, 1, :, :]
+            opd = data[:, 1, :, :]
             opd = scale(opd, 'cld_opd_dcomp', mean_std_dct)
             opd = opd.astype(np.float32)
             opd_s.append(opd)