Repository: pytorch
Issue: apex internal assert failed
Labels: module: dependency bug,triaged,module: assert failure,module: regression

## 🐛 Bug

During training in Fairseq, I get an internal assert error, which asks me to open an issue.

## To Reproduce

I run a fairseq script and see:

```
raceback (most recent call last):
  File "/private/home/jmeier/.conda/envs/fairseq-20210102-py38/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 19, in _wrap
    fn(i, *args)
  File "/private/home/jmeier/proteinseq_unsupervised_contacts/fairseq/distributed_utils.py", line 159, in distributed_main
    main(args, **kwargs)
  File "/private/home/jmeier/proteinseq_unsupervised_contacts/fairseq_cli/train.py", line 126, in main
    valid_losses, should_stop = train(args, trainer, task, epoch_itr)
  File "/private/home/jmeier/.conda/envs/fairseq-20210102-py38/lib/python3.8/contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "/private/home/jmeier/proteinseq_unsupervised_contacts/fairseq_cli/train.py", line 209, in train
    log_output = trainer.train_step(samples)
  File "/private/home/jmeier/.conda/envs/fairseq-20210102-py38/lib/python3.8/contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "/private/home/jmeier/proteinseq_unsupervised_contacts/fairseq/trainer.py", line 505, in train_step
    raise e
  File "/private/home/jmeier/proteinseq_unsupervised_contacts/fairseq/trainer.py", line 473, in train_step
    loss, sample_size_i, logging_output = self.task.train_step(
  File "/private/home/jmeier/proteinseq_unsupervised_contacts/fairseq/tasks/fairseq_task.py", line 407, in train_step
    loss, sample_size, logging_output = criterion(model, sample)
  File "/private/home/jmeier/.conda/envs/fairseq-20210102-py38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/private/home/jmeier/proteinseq_unsupervised_contacts/pretraining/masked_lm_contacts.py", line 66, in forward
    return super().forward(model, sample, reduce)
  File "/private/home/jmeier/proteinseq_unsupervised_contacts/fairseq/criterions/masked_lm.py", line 57, in forward
    loss = modules.cross_entropy(
  File "/private/home/jmeier/proteinseq_unsupervised_contacts/fairseq/modules/cross_entropy.py", line 33, in cross_entropy
    losses = xentropy.SoftmaxCrossEntropyLoss.apply(
  File "/private/home/jmeier/.conda/envs/fairseq-20210102-py38/lib/python3.8/site-packages/apex/contrib/xentropy/softmax_xentropy.py", line 7, in forward
    losses, max_log_sum_exp = xentropy_cuda.forward(
RuntimeError: labels.type().is_cuda() INTERNAL ASSERT FAILED at "apex/contrib/csrc/xentropy/interface.cpp":30, please report a bug to PyTorch. labels must be a CUDA tensor
```

## Expected behavior

Model should train fine. This error does not occur with PyTorch 1.6; only 1.8. 

## Environment

Collecting environment information...
PyTorch version: 1.7.0a0+57bffc3
Is debug build: False
CUDA used to build PyTorch: 11.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.1 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: Tesla V100-SXM2-32GB
GPU 1: Tesla V100-SXM2-32GB
GPU 2: Tesla V100-SXM2-32GB
GPU 3: Tesla V100-SXM2-32GB
GPU 4: Tesla V100-SXM2-32GB
GPU 5: Tesla V100-SXM2-32GB
GPU 6: Tesla V100-SXM2-32GB
GPU 7: Tesla V100-SXM2-32GB

Nvidia driver version: 450.80.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] msgpack-numpy==0.4.7.1
[pip3] numpy==1.19.2
[pip3] pytorch-lightning==1.0.5
[pip3] torch==1.7.0a0
[pip3] torchaudio==0.7.0a0+a853dff
[conda] blas                      1.0                         mkl
[conda] faiss-cpu                 1.6.5           py3.8_h0236df7_0_cpu    pytorch
[conda] libfaiss                  1.6.5            h2bc3f7f_0_cpu    pytorch
[conda] magma-cuda110             2.5.2                         1    pytorch
[conda] mkl                       2020.2                      256
[conda] mkl-include               2020.2                      256
[conda] mkl-service               2.3.0            py38he904b0f_0
[conda] mkl_fft                   1.2.0            py38h23d657b_0
[conda] mkl_random                1.1.1            py38h0573a6f_0
[conda] msgpack-numpy             0.4.7.1                  pypi_0    pypi
[conda] numpy                     1.19.2           py38h54aff64_0
[conda] numpy-base                1.19.2           py38hfa32c7d_0
[conda] pytorch-lightning         1.0.5                    pypi_0    pypi
[conda] torch                     1.7.0a0                  pypi_0    pypi
[conda] torchaudio                0.7.0a0+a853dff          pypi_0    pypi



cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser