Repository: pytorch
Issue: Complex backward returns NaN values
Labels: high priority,module: cuda,triaged,module: complex,complex_autograd

## 🐛 Bug

I am training this model: https://github.com/jonashaag/asteroid/blob/dccrn-recipe/asteroid/models/dccrnet.py

Tried PyTorch 1.6.x with CUDA 10.1 and PyTorch Nightly with CUDA 10.2. I can successfully train the model with CUDA 11.1 and PyTorch Nightly on another machine. I can't test CUDA 11.1 on the failing machine though.

## To Reproduce

Steps to reproduce the behavior:

Run https://github.com/jonashaag/asteroid/blob/dccrn-recipe/egs/dns_challenge/dccrn/run.sh, wait for ~100 batches -> crash

Unfortunately the autogenerated code snippet from PyTorch does not fail, and I don't have a minimal reproducing code.

```
  File "train.py", line 138, in main                                                                                                                                       [320/1659]
    trainer.fit(system)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 454, in fit
    results = self.accelerator_backend.train()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/gpu_backend.py", line 53, in train
    results = self.train_or_test()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/base_accelerator.py", line 50, in train_or_test
    results = self.trainer.train()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 501, in train
    self.train_loop.run_training_epoch()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 527, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 662, in run_training_batch
    self.trainer.hiddens
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 741, in training_step_and_backward
    result = self.training_step(split_batch, batch_idx, opt_idx, hiddens)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 302, in training_step
    training_step_output = self.trainer.accelerator_backend.training_step(args)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/gpu_backend.py", line 61, in training_step
    output = self.__training_step(args)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/gpu_backend.py", line 69, in __training_step
    output = self.trainer.model.training_step(*args)
  File "/root/asteroid/asteroid/engine/system.py", line 116, in training_step
    loss, orig_loss = self.common_step(batch, batch_nb, train=True)
  File "/root/asteroid/asteroid/engine/system.py", line 93, in common_step
    est_targets = self(inputs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/root/asteroid/asteroid/engine/system.py", line 66, in forward
    return self.model(*args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/root/asteroid/asteroid/models/base_models.py", line 248, in forward
    est_masks = self.masker(tf_rep)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/root/asteroid/asteroid/masknn/base.py", line 120, in forward
    return super().forward(x.unsqueeze(1))
  File "/root/asteroid/asteroid/masknn/base.py", line 39, in forward
    x = enc(x)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/root/asteroid/asteroid/masknn/convolutional.py", line 429, in forward
    return self.activation(self.norm(self.conv(x)))
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/root/asteroid/asteroid/complex_nn.py", line 129, in forward
    self.re_module(x.real) - self.im_module(x.imag),
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 390, in forward
    return self._conv_forward(input, self.weight)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 387, in _conv_forward
    self.padding, self.dilation, self.groups)
RuntimeError: cuDNN error: CUDNN_STATUS_EXECUTION_FAILED
```

Automated repo script (doesn't crash):

```
import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.allow_tf32 = True
data = torch.randn([20, 1, 257, 186], dtype=torch.float, device='cuda', requires_grad=True)
net = torch.nn.Conv2d(1, 32, kernel_size=[5, 2], padding=[2, 0], stride=[2, 1], dilation=[1, 1], groups=1)
net = net.cuda().float()
out = net(data)
out.backward(torch.randn_like(out))
torch.cuda.synchronize()
```

```
ConvolutionParams
    data_type = CUDNN_DATA_FLOAT
    padding = [2, 0, 0]
    stride = [2, 1, 0]
    dilation = [1, 1, 0]
    groups = 1
    deterministic = false
    allow_tf32 = true
input: TensorDescriptor 0x5588803927e0
    type = CUDNN_DATA_FLOAT
    nbDims = 4
    dimA = 20, 1, 257, 186,
    strideA = 47802, 47802, 186, 1,
output: TensorDescriptor 0x5588808e1aa0
    type = CUDNN_DATA_FLOAT
    nbDims = 4
    dimA = 20, 32, 129, 185,
    strideA = 763680, 23865, 185, 1,
weight: FilterDescriptor 0x558880392c00
    type = CUDNN_DATA_FLOAT
    tensor_format = CUDNN_TENSOR_NCHW
    nbDims = 4
    dimA = 32, 1, 5, 2,
Pointer addresses:
    input: 0x7f77c0353000
    output: 0x7f77d6000000
    weight: 0x7f7a2d600000
Forward algorithm: 1
```

When running with anomaly detection:

```
Warning: Error detected in StackBa
ckward. Traceback of forward call that caused the error:
  File "train.py", line 173, in <module>
    main(arg_dic)
  File "train.py", line 139, in main
    trainer.fit(system)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 454, in fit
    results = self.accelerator_backend.train()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/gpu_backend.py", line 53, in train
    results = self.train_or_test()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/base_accelerator.py", line 50, in train_or_test
    results = self.trainer.train()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 501, in train
    self.train_loop.run_training_epoch()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 527, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 662, in run_training_batch
    self.trainer.hiddens
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 741, in training_step_and_backward
    result = self.training_step(split_batch, batch_idx, opt_idx, hiddens)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 302, in training_step
    training_step_output = self.trainer.accelerator_backend.training_step(args)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/gpu_backend.py", line 61, in training_step
    output = self.__training_step(args)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/gpu_backend.py", line 69, in __training_step
    output = self.trainer.model.training_step(*args)
  File "/root/asteroid/asteroid/engine/system.py", line 116, in training_step
    loss, orig_loss = self.common_step(batch, batch_nb, train=True)
  File "/root/asteroid/asteroid/engine/system.py", line 93, in common_step
    est_targets = self(inputs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/root/asteroid/asteroid/engine/system.py", line 66, in forward
    return self.model(*args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/root/asteroid/asteroid/models/base_models.py", line 248, in forward
    est_masks = self.masker(tf_rep)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/root/asteroid/asteroid/masknn/base.py", line 120, in forward
    return super().forward(x.unsqueeze(1))
  File "/root/asteroid/asteroid/masknn/base.py", line 43, in forward
    x = dec(x)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/root/asteroid/asteroid/masknn/convolutional.py", line 480, in forward
    return self.activation(self.norm(self.deconv(x)))
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/root/asteroid/asteroid/complex_nn.py", line 104, in forward
    return as_torch_complex((self.re_module(x.real), self.im_module(x.imag)))
  File "/root/asteroid/asteroid/complex_nn.py", line 48, in as_torch_complex
    return torch.view_as_complex(torch.stack(x, dim=-1))

 (function _print_stack)
Traceback (most recent call last):
  File "train.py", line 173, in <module>
    main(arg_dic)
  File "train.py", line 139, in main
    trainer.fit(system)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 454, in fit
    results = self.accelerator_backend.train()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/gpu_backend.py", line 53, in train
    results = self.train_or_test()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/base_accelerator.py", line 50, in train_or_test
    results = self.trainer.train()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 501, in train
    self.train_loop.run_training_epoch()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 527, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 662, in run_training_batch
    self.trainer.hiddens
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 748, in training_step_and_backward
    self.backward(result, optimizer, opt_idx)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 278, in backward
    result.closure_loss = self.trainer.accelerator_backend.backward(result.closure_loss, optimizer, opt_idx)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/base_accelerator.py", line 90, in backward
    model_ref.backward(self, closure_loss, optimizer, opt_idx)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/core/hooks.py", line 357, in backward
    loss.backward()
  File "/opt/conda/lib/python3.7/site-packages/torch/tensor.py", line 227, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/opt/conda/lib/python3.7/site-packages/torch/autograd/__init__.py", line 138, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Function 'StackBackward' returned nan values in its 0th output.
```

## Expected behavior

Should train with no crash

## Environment

(Oops, forgot to gather this data, and the machine is already de-provisioned...)

cc @ezyang @gchanan @zou3519 @bdhirsh @heitorschueroff @anjali411 @dylanbespalko @mruberry @ngimel @ejguan