Repository: pytorch
Issue: RuntimeError: t == DeviceType::CUDAINTERNAL ASSERT FAILED when trying to calculate gradients
Labels: module: cuda,triaged

### 🐛 Describe the bug

I am working on a constrained optimization problem using the augmented lagrange method. However, after calculating the first loss term in the optimization, I received the following error after trying to calculate the gradients with `loss.backward()`.

```
/home/ihowell/Projects/validity/validity/generators/bern_vae.py:23: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  LOG2T = torch.log2(torch.tensor(1 - 2 * EPSILON))
label tensor([5])
target label tensor([6])
loss tensor([7615.2256], device='cuda:0', grad_fn=<AddBackward0>)
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/local/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/ihowell/Projects/validity/validity/contrastive/cdeepex.py", line 274, in <module>
    fire.Fire(run_xgems)
  File "/home/ihowell/Projects/validity/env/lib/python3.8/site-packages/fire/core.py", line 141, in Fire
    component_trace = _Fire(component, args, parsed_flag_args, context, name)
  File "/home/ihowell/Projects/validity/env/lib/python3.8/site-packages/fire/core.py", line 466, in _Fire
    component, remaining_args = _CallAndUpdateTrace(
  File "/home/ihowell/Projects/validity/env/lib/python3.8/site-packages/fire/core.py", line 681, in _CallAndUpdateTrace
    component = fn(*varargs, **kwargs)
  File "/home/ihowell/Projects/validity/validity/contrastive/cdeepex.py", line 222, in run_xgems
    x, x_hat = cdeepex(encode, decode, classifier, data, target_label, num_classes, writer)
  File "/home/ihowell/Projects/validity/validity/contrastive/cdeepex.py", line 90, in cdeepex
    loss.backward()
  File "/home/ihowell/Projects/validity/env/lib/python3.8/site-packages/torch/_tensor.py", line 307, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/home/ihowell/Projects/validity/env/lib/python3.8/site-packages/torch/autograd/__init__.py", line 154, in backward
    Variable._execution_engine.run_backward(
RuntimeError: t == DeviceType::CUDAINTERNAL ASSERT FAILED at "../c10/cuda/impl/CUDAGuardImpl.h":24, please report a bug to PyTorch. 
```

Here is the code I am running:

```
    x_start = x_start.cuda()
    y_true = classifier(x_start).argmax(-1)
    y_probe = y_probe.cuda()
    z_0 = encode(x_start)
    z = z_0.clone().detach()
    del_z_0 = x_start - decode(z_0)[0]
    c = 1.
    lam = torch.tensor(1.).cuda()
    mu_1 = torch.tensor(1.).cuda()
    mu_2 = torch.tensor(1.).cuda()
    beta = 1.01
    gamma = 0.24

    y_prime_idx = torch.tensor([i for i in range(num_classes)
                                if i not in [y_true, y_probe]]).cuda()

    def loss_c(z, lam, mu_1, mu_2, c):
        l = (z - z_0).norm(dim=-1)
        l += lam * h(z, y_true, y_probe)
        l += c / 2 * h(z, y_true, y_probe).norm(dim=-1)
        l += 1 / (2 * c) * (torch.maximum(ZERO, mu_1 + c * h(z, y_true, y_prime_idx)).square()
                            - mu_1.square()).sum()
        l += 1 / (2 * c) * (torch.maximum(
            ZERO, mu_2 + c * h(z, y_probe, y_prime_idx)).square() - mu_2.square()).sum()
        return l

    def h(z, y_1, y_2):
        img = decode(z)[0] + del_z_0
        logits = classifier(img)
        return (logits.index_select(1, y_1) - logits.index_select(1, y_2)).squeeze(-1)

    for i in range(outer_iters):
        old_z = z.clone().detach()
        early_stopping = EarlyStopping(patience=10)
        optimizer = optim.Adam([z], weight_decay=1e-5)
        for j in range(inner_iters):
            optimizer.zero_grad()
            loss = loss_c(z, lam, mu_1, mu_2, c)
            print('loss', loss)
            loss.backward()
            optimizer.step()
            if early_stopping(loss):
                break

        lam += c * h(z, y_true, y_probe)
        mu_1 = torch.maximum(ZERO, mu_1 + c * h(z, y_true, y_prime_idx))
        mu_2 = torch.maximum(ZERO, mu_2 + c * h(z, y_true, y_prime_idx))

        if h(z, y_true, y_probe) > gamma * h(old_z, y_true, y_probe):
            c = c * beta

        writer.add_scalar('cdeepex/loss', loss, i)
        img = torch.cat([decode(z_0)[0], decode(z)[0]], 3)[0]
        writer.add_image('cdeepex/example', torch.tensor(img), i)
        print(f'Min loss: {loss}')

    return decode(z_0)[0], decode(z)[0]
```

### Versions

```
PyTorch version: 1.10.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.21.2
Libc version: glibc-2.27

Python version: 3.8.6 (default, Aug 25 2021, 16:01:17)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-91-generic-x86_64-with-glibc2.2.5
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1080
Nvidia driver version: 470.82.00
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.4
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.21.3
[pip3] torch==1.10.0
[pip3] torchvision==0.9.0
[conda] Could not collect
```

cc @ngimel