Error when calculate tokens_id and Mask LM
## 🐛 Bug

<!-- Important information -->

Model I am using (DistilBert):

Language I am using the model on (English):

The problem arise when using: Distiller.prepare_batch( ) 
Error when token_ids is masked by mask LM matrix
*  the official example scripts:
_token_ids_real = token_ids[pred_mask] 
*  my own modified scripts: 
 _token_ids_real=torch.mul(token_ids, pred_mask)

The tasks I am working on is:
* [GLUE ] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details)

## To Reproduce

Steps to reproduce the behavior:
1. pred_mask is matrix with 0,1. 
Operation token_ids[pred_mask] seems to make some same matrix, instead of masking token_ids

<!-- If you have a code sample, error messages, stack traces, please provide it here as well. -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

* OS: Win10
* Python version: 3.6
* PyTorch version: 1.1
* PyTorch Transformers version (or branch): 2.0/0
* Using GPU ? Yes
* Distributed of parallel setup ? No
* Any other relevant information: 

## Additional context

<!-- Add any other context about the problem here. -->
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
