From c7e26ea1887a9353f387e86571ab2065bca16a61 Mon Sep 17 00:00:00 2001
From: Kirill Kouzoubov <kirill888@gmail.com>
Date: Sun, 2 Jun 2024 14:50:20 +1000
Subject: [PATCH] sqme: load_chunks fix

---
 odc/loader/_builder.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/odc/loader/_builder.py b/odc/loader/_builder.py
index 605f409..7b75443 100644
--- a/odc/loader/_builder.py
+++ b/odc/loader/_builder.py
@@ -545,14 +545,14 @@ def load_tasks(
     instances for every possible time, y, x, bins, including empty ones.
     """
     # pylint: disable=too-many-locals
+    extra_dims = extra_dims or {}
+    chunks = chunks or {}
+
     if nt is None:
         nt = max(t for t, _, _ in tyx_bins) + 1
 
-    if extra_dims is None:
-        extra_dims = {}
-    if chunks is None:
-        chunks = {}
-
+    chunks = {**chunks}
+    chunks.update(zip(["y", "x"], gbt.chunk_shape((0, 0)).yx))
     base_shape = (nt, *gbt.base.shape.yx)
 
     if bands is None:
@@ -682,6 +682,8 @@ def _do_one(task: LoadChunkTask) -> Tuple[str, int, int, int]:
         nt=nt,
         extra_dims=template.extra_dims_full(),
     )
+    tasks = list(tasks)
+    assert len(tasks) == total_tasks
 
     _work = pmap(_do_one, tasks, pool)