diff --git a/km3io/offline.py b/km3io/offline.py
index 5d551666ad6e847f55515c2e035e53264321bf6d..8065bcca6e4b0d35fd869e9ce6b0158f94df03d1 100644
--- a/km3io/offline.py
+++ b/km3io/offline.py
@@ -568,8 +568,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)
@@ -690,8 +692,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