[Python] Configuring a progress bar while training for Deep Learning

I have this tiny training function upcycled from a tutorial.

def train(epoch, tokenizer, model, device, loader, optimizer):
model.train()
with tqdm.tqdm(loader, unit="batch") as tepoch:
for _,data in enumerate(loader, 0):
y = data['target_ids'].to(device, dtype = torch.long)
y_ids = y[:, :-1].contiguous()
lm_labels = y[:, 1:].clone().detach()
lm_labels[y[:, 1:] == tokenizer.pad_token_id] = -100
ids = data['source_ids'].to(device, dtype = torch.long)
mask = data['source_mask'].to(device, dtype = torch.long)

outputs = model(input_ids = ids, attention_mask = mask, decoder_input_ids=y_ids, labels=lm_labels)
loss = outputs[0]

tepoch.set_description(f"Epoch {epoch}")
tepoch.set_postfix(loss=loss.item())

if _%10 == 0:
wandb.log({"Training Loss": loss.item()})

if _%1000==0:
print(f'Epoch: {epoch}, Loss: {loss.item()}')

optimizer.zero_grad()
loss.backward()
optimizer.step()
# xm.optimizer_step(optimizer)
# xm.mark_step()


The function trains fine, the problem is that I can't seem to make the progress bar work correctly. I played around with it, but haven't found a configuration that correctly updates the loss and tells me how much time is left. What might I be doing wrong?

Continue reading...