diff --git a/dl1_data_handler/writer.py b/dl1_data_handler/writer.py
index 93b8d122..f4c13dc7 100644
--- a/dl1_data_handler/writer.py
+++ b/dl1_data_handler/writer.py
@@ -992,6 +992,9 @@ def gain_selection(self, waveform, image, peakpos, cam_id, threshold):
         self.gain_selector.thresholds[cam_id] = threshold
 
         waveform, gain_mask = self.gain_selector.select_gains(cam_id, waveform)
+        if gain_mask.ndim == 1:
+            gain_mask = gain_mask[:, None]
+
         signal_mask = gain_mask.max(axis=1)
 
         combined_image = image[0].copy()