FusedLayerNorm leads to RuntimeError: CUDA error: no kernel image is available for execution on the device

After a GPU tensor goes through FusedLayerNorm, the next time that memory is accessed I get a RuntimeError: CUDA error: no kernel image is available for execution on the device.

To reproduce:

import torch
from apex.normalization import FusedLayerNorm

norm = FusedLayerNorm(16)
device = torch.device('cuda:' + str(0))
norm = norm.to(device)

x = torch.randn(3, 4, 16)
x = x.to(device)
attended = norm(x)
print(x)

Other operations on attended or x will also raise the error. However, if I move x to the CPU, I can then proceed to use it without any problems.

I'm running this on an AWS p3.2xlarge instance based on the AWS Deep Learning AMI (Ubuntu 18.04) Version 25.0. We've updated pytorch to 1.3.0, and installed GPUtil, Apex, and gpustat using the following commands:

source activate pytorch_p36

# Update to the latest PyTorch 1.3 (but not CUDA 10.0 instead of 10.1, because the AMI/env doesn't have it installed)
conda install pytorch==1.3.0 torchvision==0.4.1 cudatoolkit=10.0 -c pytorch -y

# Install GPUtil
pip install GPUtil

# Install NVIDIA Apex
git clone https://github.com/NVIDIA/apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

# Install gpustat
pip install gpustat

Doing the same thing on an aws p2.xlarge instance with the same changes to the environment does not cause the error.

Start with the provided FusedLayerNorm reproduction on the AWS p3.2xlarge environment, then compare it with the working p2.xlarge setup and their PyTorch, CUDA toolkit, and Apex builds. Done means GPU access after FusedLayerNorm no longer raises the kernel-image RuntimeError for the reproduced operations.