fix(flagging): redistribute back to frequency after generating baseli…

…ne mask
radiocosmology · Dec 4, 2023 · a66283b · a66283b
1 parent 1fb52bf
commit a66283b
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/draco/analysis/flagging.py b/draco/analysis/flagging.py
@@ -933,6 +933,10 @@ def process(
         # Save mask to output container
         mask_cont.mask[:] = mpiarray.MPIArray.wrap(local_mask, axis=1)
 
+        # Distribute back across frequency
+        mask_cont.redistribute("freq")
+        stream.redistribute("freq")
+
         return mask_cont