tqdm_train_loop.py

from tqdm.auto import trange, tqdm
def tqdm_train_loop(train_fn, num_epochs, train_loader, val_fn=None):
    # each epoch
    with trange(num_epochs, unit="epoch") as epoch_t:
        for epoch in epoch_t:
            # each batch
            with tqdm(train_loader, leave=False, unit='batch', postfix='') as batch_t:
                for batch_idx, (features, targets) in enumerate(batch_t):
                    cost = train_fn(features, labels)
                    batch_t.postfix = f"cost {cost.item():.2f}"
                    batch_t.update()    
            # compute validation loss, if possible
            if val_fn:
                val_cost = val_fn(features, labels)                                                    
                epoch_t.postfix = f"train {cost.item():.2f}, val {val_cost.item():.2f} "
            else:
                epoch_t.postfix = f"train {cost.item():.2f}"                
            epoch_t.update()