[examples] fix no grad in second pruning in run_bertology
Author: TobiasLee
Pull request: yes

the `new_head_mask` index assignment operation makes it become a non-leaf node in the following gradient computation, resulting in grad is None bug as mentioned in  #3895 


Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/4479?src=pr&el=h1) Report
> :exclamation: No coverage uploaded for pull request base (`master@18d233d`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference#section-missing-base-commit).
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/4479/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/4479?src=pr&el=tree)

```diff
@@            Coverage Diff            @@
##             master    #4479   +/-   ##
=========================================
  Coverage          ?   78.20%           
=========================================
  Files             ?      120           
  Lines             ?    20083           
  Branches          ?        0           
=========================================
  Hits              ?    15705           
  Misses            ?     4378           
  Partials          ?        0           
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/4479?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/modeling\_tf\_albert.py](https://codecov.io/gh/huggingface/transformers/pull/4479/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9hbGJlcnQucHk=) | `78.74% <0.00%> (ø)` | |
| [src/transformers/configuration\_marian.py](https://codecov.io/gh/huggingface/transformers/pull/4479/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX21hcmlhbi5weQ==) | `100.00% <0.00%> (ø)` | |
| [src/transformers/data/data\_collator.py](https://codecov.io/gh/huggingface/transformers/pull/4479/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9kYXRhL2RhdGFfY29sbGF0b3IucHk=) | `89.23% <0.00%> (ø)` | |
| [src/transformers/configuration\_transfo\_xl.py](https://codecov.io/gh/huggingface/transformers/pull/4479/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX3RyYW5zZm9feGwucHk=) | `92.45% <0.00%> (ø)` | |
| [src/transformers/tokenization\_transfo\_xl.py](https://codecov.io/gh/huggingface/transformers/pull/4479/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdHJhbnNmb194bC5weQ==) | `40.67% <0.00%> (ø)` | |
| [src/transformers/data/processors/xnli.py](https://codecov.io/gh/huggingface/transformers/pull/4479/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9kYXRhL3Byb2Nlc3NvcnMveG5saS5weQ==) | `29.54% <0.00%> (ø)` | |
| [src/transformers/hf\_argparser.py](https://codecov.io/gh/huggingface/transformers/pull/4479/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9oZl9hcmdwYXJzZXIucHk=) | `61.97% <0.00%> (ø)` | |
| [src/transformers/tokenization\_bart.py](https://codecov.io/gh/huggingface/transformers/pull/4479/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fYmFydC5weQ==) | `100.00% <0.00%> (ø)` | |
| [src/transformers/tokenization\_flaubert.py](https://codecov.io/gh/huggingface/transformers/pull/4479/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fZmxhdWJlcnQucHk=) | `40.42% <0.00%> (ø)` | |
| [src/transformers/data/processors/\_\_init\_\_.py](https://codecov.io/gh/huggingface/transformers/pull/4479/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9kYXRhL3Byb2Nlc3NvcnMvX19pbml0X18ucHk=) | `100.00% <0.00%> (ø)` | |
| ... and [110 more](https://codecov.io/gh/huggingface/transformers/pull/4479/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/4479?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/4479?src=pr&el=footer). Last update [18d233d...b3c4f81](https://codecov.io/gh/huggingface/transformers/pull/4479?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


This looks reasonable to me. Thanks for looking into it.

can you run 

```
pip uninstall -y isort black
pip install -e .[quality]
make style
```
?

Thanks!

Thanks for noting how to run code reformatting! 

Thanks!