progress bar

src/edm.py

@@ -8,6 +8,8 @@ from src.egnn import Dynamics
 from src.noise import GammaNetwork, PredefinedNoiseSchedule
 from typing import Union
 
+from tqdm import tqdm
+
 from pdb import set_trace
 
 
@@ -565,7 +567,7 @@ class InpaintingEDM(EDM):
         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