diff --git a/km3io/offline.py b/km3io/offline.py
index c71b49abf066a8c99ac01b9514eb4739ef3a39a9..ba4d0e5529f8916da81274673562f7d6b4bf5dc0 100644
--- a/km3io/offline.py
+++ b/km3io/offline.py
@@ -569,8 +569,10 @@ class OfflineReader:
                     str(stages)))
         else:
             fit_data = np.array([
-                i[k] for i, j, k in zip(fit_info, rec_stages[:, 0],
-                                        rec_stages[:, 1]) if k is not None
+                i[k]
+                for i, j, k in zip(fit_info, rec_stages[:, 0], rec_stages[:,
+                                                                          1])
+                if k is not None
             ])
             rec_array = np.core.records.fromarrays(fit_data.transpose(),
                                                    names=keys)
@@ -691,8 +693,8 @@ class OfflineReader:
             for key in keys:
                 lazy_d[key] = np.array([
                     i[k] for i, k in zip(
-                        getattr(self.tracks, key)[mask], rec_stages[:, 1]
-                        [mask])
+                        getattr(self.tracks, key)[mask], rec_stages[:,
+                                                                    1][mask])
                 ])
 
         return lazy_d