Optimize Progress Bar Implementation in train.py

[heouchen]

The code in train.py(line:123-124)

with tqdm(total=len(train_set), desc=f'Epoch {epoch + 1} / {epochs}',
          unit='img') as pbar:

showed be modify to:

with tqdm(total=len(train_loader), desc=f'Epoch {epoch + 1} / {epochs}',
          unit='batch') as pbar:

[mahmoudnafifi]

Thanks! Could you please create a PR with the fix?