Skip to content

GitLab

Commit 335f57ba authored by thomwolf's avatar thomwolf
Browse files

only on main process

parent 326944d6
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 3 deletions
examples/run_squad.py
View file @ 335f57ba
...@@ -917,7 +917,8 @@ def main(): ...@@ -917,7 +917,8 @@ def main():
model = torch.nn.DataParallel(model) model = torch.nn.DataParallel(model)
if args.do_train: if args.do_train:
writer = SummaryWriter() if args.local_rank in [-1, 0]:
writer = SummaryWriter()
# Prepare data loader # Prepare data loader
train_examples = read_squad_examples( train_examples = read_squad_examples(
input_file=args.train_file, is_training=True, version_2_with_negative=args.version_2_with_negative) input_file=args.train_file, is_training=True, version_2_with_negative=args.version_2_with_negative)
...@@ -1016,8 +1017,9 @@ def main(): ...@@ -1016,8 +1017,9 @@ def main():
else: else:
loss.backward() loss.backward()
if (step + 1) % args.gradient_accumulation_steps == 0: if (step + 1) % args.gradient_accumulation_steps == 0:
writer.add_scalar('lr', optimizer.get_lr()[0], global_step) if args.local_rank in [-1, 0]:
writer.add_scalar('loss', loss.item(), global_step) writer.add_scalar('lr', optimizer.get_lr()[0], global_step)
writer.add_scalar('loss', loss.item(), global_step)
if args.fp16: if args.fp16:
# modify learning rate with special warm up BERT uses # modify learning rate with special warm up BERT uses
# if args.fp16 is False, BertAdam is used and handles this automatically # if args.fp16 is False, BertAdam is used and handles this automatically
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment