

TITLE
KeyError: loss when pretraining using BertForPreTraining

COMMENTS
1

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
### System Info

```shell
- `transformers` version: 4.19.0.dev0
- Platform: Linux-5.13.0-40-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.5.1
- PyTorch version (GPU?): 1.11.0+cu102 (False)
- Tensorflow version (GPU?): 2.7.0 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>
```


### Who can help?

@sgugger 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

==========================================

My dataset is in `.txt` format, and it look like this:
```
Sentence 1
Sentence 2
Sentence 3
Sentence 4

Sentence A
Sentence B
Sentence C
Sentence D
Sentence E

Sentence a
Sentence b
Sentence c
```

### Reproduction

1. Tokenize my own dataset using WordLevel Tokenizer
2. Do post pre-processing
3. Train the tokenizer
4. Load Dataset using LineByLineTextDataset
5. Define the configuration of BERT model using BertConfig
6. Create BertForPreTraining model
7. Define Data Collator using DataCollatorForLanguageModeling
8. Initialize Trainer
9. Do pre-training

=======================================================

Below are the code snippets from step 1 to 9.

### 1) Tokenize Dataset Using WordLevel Tokenizer
```
tokenizer = Tokenizer(WordLevel(unk_token="[UNK]"))

trainer = WordLevelTrainer(vocab_size=52_000, 
                           min_frequency=1, 
                           special_tokens=["[UNK]", 
                                           "[CLS]", 
                                           "[SEP]", 
                                           "[PAD]", 
                                           "[MASK]"])
tokenizer.pre_tokenizer = Whitespace()
```

### 2) Do Post Pre-Processing
```
tokenizer.post_processor = TemplateProcessing(
    single="[CLS] $A [SEP]",
    pair="[CLS] $A [SEP] $B:1 [SEP]:1",
    special_tokens=[
        ("[CLS]", 1),
        ("[SEP]", 2),
    ],
)
```

### 3) Train Tokenizer
```
path = [str(x) for x in Path("path_to_text_corpus").glob("**/*.txt")]
tokenizer.train(path, trainer)
tokenizer.save("path_to_trained_tokenizer.json")
```

### 4) Load Dataset
```
dataset = LineByLineTextDataset(
    tokenizer=tokenizer,
    file_path="path_to_corpus.txt",
    block_size=128,
)
```

### 5) Define BERT Configuration
```
config = BertConfig(
    vocab_size=50000,
    hidden_size=768,
    num_hidden_layers=12, 
    num_attention_heads=12, 
    intermediate_size=3072, hidden_act='gelu', 
    hidden_dropout_prob=0.1, 
    attention_probs_dropout_prob=0.1, 
    max_position_embeddings=512, 
    type_vocab_size=1, 
    initializer_range=0.02, 
    layer_norm_eps=1e-12, 
    pad_token_id=3, 
    gradient_checkpointing=False,
)
```

### 6) Create BERT Model for Pretraining
```
model = BertForPreTraining(config=config)
```

### 7) Define Data Collator
```
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, 
    mlm=True, 
    mlm_probability=0.15,
)
```

### 8) Initialize Trainer
```
training_args = TrainingArguments(
    output_dir="path_to_pretrained_model",
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=4,
    save_steps=10000,
    save_total_limit=2,
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=dataset,
)
```

### 9) Do Pre-Training
```
trainer.train()
```

### Expected behavior

I want to pretrain BERT like model with NSP and MLM, but when i run `trainer.train`, i got this error:

```shell


---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<timed eval> in <module>

~/env/lib/python3.8/site-packages/transformers/trainer.py in train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1421                         tr_loss_step = self.training_step(model, inputs)
   1422                 else:
-> 1423                     tr_loss_step = self.training_step(model, inputs)
   1424 
   1425                 if (

~/env/lib/python3.8/site-packages/transformers/trainer.py in training_step(self, model, inputs)
   2010 
   2011         with self.autocast_smart_context_manager():
-> 2012             loss = self.compute_loss(model, inputs)
   2013 
   2014         if self.args.n_gpu > 1:

~/env/lib/python3.8/site-packages/transformers/trainer.py in compute_loss(self, model, inputs, return_outputs)
   2052         else:
   2053             # We don't use .loss here since the model may return tuples instead of ModelOutput.
-> 2054             loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]
   2055 
   2056         return (loss, outputs) if return_outputs else loss

~/env/lib/python3.8/site-packages/transformers/utils/generic.py in __getitem__(self, k)
    217         if isinstance(k, str):
    218             inner_dict = {k: v for (k, v) in self.items()}
--> 219             return inner_dict[k]
    220         else:
    221             return self.to_tuple()[k]

KeyError: 'loss'
```
```
