Better tqdm for tps_baseline

Author: plainerman

tps_baseline.py

@@ -365,6 +365,6 @@ def step_langevin_backward(_x, _v, _key):
     plt.show()
     plt.clf()
 
-    for i, path in tqdm(enumerate(paths)):
+    for i, path in tqdm(enumerate(paths), desc='Saving trajectories', total=len(paths)):
         save_trajectory(mdtraj_topology, jnp.array([kabsch_align(p.reshape(-1, 3), B.reshape(-1, 3))[0] for p in path]),
                         f'{savedir}/trajectory_{i}.pdb')