Changing the loss function in BERT
# ❓ Questions & Help

Hello, 

I'm trying to replace the _CrossEntropyLoss_ with the _KLDivLoss_ for a sentence classification task using some golden logits and the logits from the BERT model.

```
golden_logits = …

outputs = model(**inputs)

# the logits from BERT
logits = outputs[1]

loss_fct = KLDivLoss(reduction='sum')
loss = loss_fct(F.log_softmax(logits, dim=-1), F.softmax(golden_logits, dim=-1))
```

Am I doing this correctly ? 
This is a very general PyTorch question. Please post it on Stack Overflow. https://stackoverflow.com/

About the BERT part: `outputs` is a tuple. `[0]`  is the last hidden state (sequence output), `[1]` is the pooled output.

