Repository: pytorch
Issue: 16-bit input + AMP breaks AdaptiveLogSoftmaxWithLoss
Labels: triaged,module: amp (automated mixed precision)

## 🐛 Bug

When feeding [AdaptiveLogSoftmaxWithLoss](https://pytorch.org/docs/master/generated/torch.nn.AdaptiveLogSoftmaxWithLoss.html) 16-bit input with amp.autocast enabled it will fail with:
`RuntimeError: Expected object of scalar type Half but got scalar type Float for argument #4 'source' in call to _th_index_copy_`

This happens because AdaptiveLogSoftmax first creates an output tensor with the type of input (float16):
https://github.com/pytorch/pytorch/blob/5f692a67dba3f8f01e3c4b1cb52a8624e1056874/torch/nn/modules/adaptive.py#L173

and later trys to collect the clusters' softmax results within said tensor **in-place**, but softmax outputs float32 in amp:
https://github.com/pytorch/pytorch/blob/5f692a67dba3f8f01e3c4b1cb52a8624e1056874/torch/nn/modules/adaptive.py#L200-L202

## To Reproduce

Steps to reproduce the behavior:

Have this script:
```python
import torch
# 32 bit precision works
adaptive16b = torch.nn.AdaptiveLogSoftmaxWithLoss(5, 20, [5]).to(dtype=torch.float32, device='cuda')
float16b = torch.randn([5, 5], dtype=torch.float32, device='cuda')  # - dtype: float16
float16b_adaptive = adaptive16b(float16b, torch.tensor([1,2,3,4,5], device='cuda'))

# 16 bit precision works
adaptive16b = torch.nn.AdaptiveLogSoftmaxWithLoss(5, 20, [5]).to(dtype=torch.float16, device='cuda')
float16b = torch.randn([5, 5], dtype=torch.float16, device='cuda')  # - dtype: float16
float16b_adaptive = adaptive16b(float16b, torch.tensor([1,2,3,4,5], device='cuda'))

# 16 bit with amp breaks
with torch.cuda.amp.autocast():
    adaptive16b = torch.nn.AdaptiveLogSoftmaxWithLoss(5, 20, [5]).to(dtype=torch.float16, device='cuda')
    float16b = torch.randn([5, 5], dtype=torch.float16, device='cuda')  # - dtype: float16
    float16b_adaptive = adaptive16b(float16b, torch.tensor([1,2,3,4,5], device='cuda'))
```

**It's important here, that the forward() reaches within the tail clusters, as there is no problem if all targets are within the head.**

<details>
  <summary>Complete stack trace in here</summary>

```
Traceback (most recent call last):
  File "C:/Users/anarc/AppData/Roaming/JetBrains/PyCharm2020.2/scratches/torch/scratch.py", line 49, in <module>
    float16b_adaptive = adaptive16b(float16b, torch.tensor([1,2,3,4,5], device='cuda'))
  File "C:\tools\miniconda3\envs\ElmoOnTorch\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "C:\tools\miniconda3\envs\ElmoOnTorch\lib\site-packages\torch\nn\modules\adaptive.py", line 202, in forward
    output.index_copy_(0, row_indices, local_logprob.squeeze(1))
RuntimeError: Expected object of scalar type Half but got scalar type Float for argument #4 'source' in call to _th_index_copy_
```
</details>

## Expected behavior

Expected AdaptiveLogSoftmaxWithLoss not to produce an error and return a float32 tensor.

## Environment

PyTorch version: 1.6.0
Is debug build: No
CUDA used to build PyTorch: 10.2

<details>
  <summary>Rest of the environment in here</summary>

OS: Microsoft Windows 10 Education
GCC version: Could not collect
CMake version: Could not collect

Python version: 3.8
Is CUDA available: Yes
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: GeForce GTX 1050
Nvidia driver version: 452.06
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.19.1
[pip3] pytorch-lightning==0.9.0
[pip3] torch==1.6.0
[conda] blas                      2.16                        mkl    conda-forge
[conda] cudatoolkit               10.2.89              h74a9793_1  
[conda] libblas                   3.8.0                    16_mkl    conda-forge
[conda] libcblas                  3.8.0                    16_mkl    conda-forge
[conda] liblapack                 3.8.0                    16_mkl    conda-forge
[conda] liblapacke                3.8.0                    16_mkl    conda-forge
[conda] mkl                       2020.1                      216  
[conda] numpy                     1.19.1           py38h72c728b_0    conda-forge
[conda] pytorch                   1.6.0           py3.8_cuda102_cudnn7_0    pytorch
[conda] pytorch-lightning         0.9.0                    pypi_0    pypi
</details>

## Additional context

The module's log_prob() and predict() both work, since neither use the in-place operator.


cc @mcarilli