diff --git a/src/calng/JungfrauCorrection.py b/src/calng/JungfrauCorrection.py
index 553f86dba101adb2876e8a906e084cd0839a9bff..5aa8bb5cb22eb402ae153e32a86b8e3a18f00acb 100644
--- a/src/calng/JungfrauCorrection.py
+++ b/src/calng/JungfrauCorrection.py
@@ -346,8 +346,11 @@ class JungfrauCorrection(BaseCorrection):
         if len(cell_table.shape) == 0:
             cell_table = cell_table[np.newaxis]
         try:
+            gain_map = data_hash.get("data.gain")
+            if self.unsafe_get("dataFormat.overrideInputAxisOrder"):
+                gain_map.shape = self.input_data_shape
             self.kernel_runner.load_data(
-                image_data, data_hash.get("data.gain"), cell_table
+                image_data, gain_map, cell_table
             )
         except ValueError as e:
             self.log_status_warn(f"Failed to load data: {e}")