RuntimeError: CUDA error: the provided PTX was compiled with an unsupported toolchain.

We are attempting to train a neural net using Colab’s GPU offerings. When getting paired with an A100 GPU, we experience the following issue:

RuntimeError: CUDA error: the provided PTX was compiled with an unsupported toolchain.
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

To verify that it wasn’t something to do with the actual program we were running, we ran the same program on a GCP instance and with a Colab V100. In both cases, the program ran fine without issue/as expected.

To further verify that the program wasn’t the issue, we ran the following code snippet:

import torch
torch.tensor([0,1],device="cuda")

And had the same issue.

Is anyone else experiencing this issue?

[craigcitro]

For anyone here for JAX issues: a quick workaround until we can get a fix out is to run

!pip install https://storage.googleapis.com/jax-releases/cuda111/jaxlib-0.1.71+cuda111-cp37-none-manylinux2010_x86_64.whl

and then restart your kernel (ctrl-m . or Runtime -> Restart runtime).

[x4nth055]

Fun fact: Now I’m no longer getting an A100 after the fix.