I reran the code using the latest PyTorch. I am seeing the following error message. This time it happens with the first batch of the first epoch.
Traceback (most recent call last):
  File "run_E2E_EL_RE.py", line 962, in <module>
    main()
  File "run_E2E_EL_RE.py", line 913, in main
    global_step, tr_loss = train(args, model, tokenizer)
  File "run_E2E_EL_RE.py", line 246, in train
    ner_loss, last_hidden_states = model.forward(**ner_inputs)
  File "/home/rajarshi_kingsaint_bhowmik/anaconda3/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 167, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/home/rajarshi_kingsaint_bhowmik/anaconda3/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 177, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/home/rajarshi_kingsaint_bhowmik/anaconda3/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 86, in parallel_apply
    output.reraise()
  File "/home/rajarshi_kingsaint_bhowmik/anaconda3/lib/python3.8/site-packages/torch/_utils.py", line 429, in reraise
    raise self.exc_type(msg)
StopIteration: Caught StopIteration in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/home/rajarshi_kingsaint_bhowmik/anaconda3/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
    output = module(*input, **kwargs)
  File "/home/rajarshi_kingsaint_bhowmik/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/rajarshi_kingsaint_bhowmik/E2E-EL-RE/modeling_E2E_EL_RE.py", line 65, in forward
    mention_outputs = self.bert_mention.bert(
  File "/home/rajarshi_kingsaint_bhowmik/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/rajarshi_kingsaint_bhowmik/E2E-EL-RE/modeling_bert.py", line 758, in forward
    extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
StopIteration