Update.

diff --git a/reasoning.py b/reasoning.py
index 92699e8..003806a 100755 (executable)
--- a/reasoning.py
+++ b/reasoning.py
@@ -307,7 +307,13 @@ class Reasoning(problem.Problem):
         prompts = torch.zeros(nb, self.height, self.width * 3, dtype=torch.int64)
         answers = torch.zeros(nb, self.height, self.width, dtype=torch.int64)
         w = self.width
-        for prompt, answer in zip(prompts, answers):
+
+        for prompt, answer in tqdm.tqdm(
+            zip(prompts, answers),
+            dynamic_ncols=True,
+            desc="world generation",
+            total=prompts.size(0),
+        ):
             A = prompt[:, 0 * w : 1 * w]
             f_A = prompt[:, 1 * w : 2 * w]
             B = prompt[:, 2 * w : 3 * w]