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=&quot;batch&quot;) 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&quot;Epoch {epoch}&quot;)
      tepoch.set_postfix(loss=loss.item())
      
      if _%10 == 0:
          wandb.log({&quot;Training Loss&quot;: 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.
Does anyone have any pointers on what I might be doing wrong?
Thanks in advance!

In case anyone else has run in my same issue, thanks to the previous response I was able to configure the progress bar as I wanted with just a little tweak of what I was doing before:
def train(epoch, tokenizer, model, device, loader, optimizer):
  model.train()    
  for _,data in tqdm(enumerate(loader, 0), unit=&quot;batch&quot;, total=len(loader)):

everything stays the same, and now I have a progress bar showing percentage and loss. I prefer this solution because it allows me to keep the other logging functions I had without further changes.
