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]