🐛 Fix concatenation

Author: shaneahmed

tiatoolbox/models/engine/engine_abc.py

@@ -646,7 +646,7 @@ def save_predictions(
                 keys_to_compute = [k for k in keys_to_compute if k not in zarr_group]
             write_tasks = []
             for key in keys_to_compute:
-                dask_array = processed_predictions[key]
+                dask_array = processed_predictions[key].rechunk("auto")
                 task = dask_array.to_zarr(
                     url=save_path,
                     component=key,

tiatoolbox/models/engine/semantic_segmentor.py

@@ -1040,7 +1040,9 @@ def merge_vertical_chunkwise(
 
             probabilities_da = None
 
-        curr_chunk, curr_count = next_chunk, next_count
+        if next_chunk is not None:
+            curr_chunk, curr_count = next_chunk[overlap:], next_count[overlap:]
+
         if i + 2 < num_chunks:
             next_chunk = canvas.blocks[i + 2, 0].compute()
             next_count = count.blocks[i + 2, 0].compute()