Switch to tqdm.autonotebook, refactor tqdm call

radionets-project · Jan 20, 2025 · 5c1afce · 5c1afce
1 parent 9756ec8
commit 5c1afce
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/pyvisgen/simulation/visibility.py b/pyvisgen/simulation/visibility.py
@@ -2,7 +2,7 @@
 
 import torch
 import toma
-from tqdm import tqdm
+from tqdm.autonotebook import tqdm
 
 import pyvisgen.simulation.scan as scan
 
@@ -176,9 +176,13 @@ def _batch_loop(
         Visibilities dataclass object.
     """
     batches = torch.arange(bas[:].shape[1]).split(batch_size)
-
-    if show_progress:
-        batches = tqdm(batches)
+    batches = tqdm(
+        batches,
+        position=0,
+        disable=not show_progress,
+        desc="Computing visibilities",
+        postfix=f"Batchsize: {batch_size}",
+    )
 
     for p in batches:
         bas_p = bas[:][:, p]