Clean up tqdm progress bars during training

chemprop · Aug 25, 2020 · ae38722 · ae38722
1 parent 76e9b68
commit ae38722
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/chemprop/train/predict.py b/chemprop/train/predict.py
@@ -24,7 +24,7 @@ def predict(model: MoleculeModel,
 
     preds = []
 
-    for batch in tqdm(data_loader, disable=disable_progress_bar):
+    for batch in tqdm(data_loader, disable=disable_progress_bar, leave=False):
         # Prepare batch
         batch: MoleculeDataset
         mol_batch, features_batch = batch.batch_graph(), batch.features()

diff --git a/chemprop/train/train.py b/chemprop/train/train.py
@@ -42,7 +42,7 @@ def train(model: MoleculeModel,
     model.train()
     loss_sum, iter_count = 0, 0
 
-    for batch in tqdm(data_loader, total=len(data_loader)):
+    for batch in tqdm(data_loader, total=len(data_loader), leave=False):
         # Prepare batch
         batch: MoleculeDataset
         mol_batch, features_batch, target_batch = batch.batch_graph(), batch.features(), batch.targets()