progress bar

src/edm.py

@@ -145,7 +145,7 @@ class EDM(torch.nn.Module):
         chain = torch.zeros((keep_frames,) + z.size(), device=z.device)
 
         # Sample p(z_s | z_t)
-        for s in reversed(range(0, self.T)):
+        for s in tqdm(reversed(range(0, self.T)), total=self.T):
             s_array = torch.full((n_samples, 1), fill_value=s, device=z.device)
             t_array = s_array + 1
             s_array = s_array / self.T