[FIX] not training when epoch is small (#3006)

* solving bug where for small epochs and large gradient_accumulation_steps we never train

* black formatting

* no need to change these files
This commit is contained in:
mataney
2020-03-19 17:21:21 +02:00
committed by GitHub
parent ad7233fc01
commit c44a17db1b

View File

@@ -233,7 +233,11 @@ def train(args, train_dataset, model, tokenizer):
loss.backward()
tr_loss += loss.item()
if (step + 1) % args.gradient_accumulation_steps == 0:
if (step + 1) % args.gradient_accumulation_steps == 0 or (
# last step in epoch but step is always smaller than gradient_accumulation_steps
len(epoch_iterator) <= args.gradient_accumulation_steps
and (step + 1) == len(epoch_iterator)
):
if args.fp16:
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
else: