Repository: pytorch
Issue: cublasSgemmStridedBatched failure when calling grad of grad
Labels: needs reproduction,module: crash,module: cuda,triaged

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
CUDA Error when calling autograd twice on same graph.

## To Reproduce
I haven't been able to reproduce. If I understand better where 'cublasSgemmStridedBatched' is used maybe I can force the error. I have copied both the attempt to reproduce and the stack trace from the original error below. 

Current attempt to reproduce does not fail: 

```
import torch

def grads(x):
    m, n, p = x.shape
    n_flat = n * p
    
    x = x.reshape(m, n_flat)
    x_flat = [x[..., i].requires_grad_(True) for i in range(n_flat)]
    x_prime = torch.stack(x_flat, -1)
    x = x_prime.view(m, n, p)
    y = function(x)
    
    dy_dx = torch.autograd.grad(y.sum(), x_prime, create_graph=True, retain_graph=True)[0]
    dy_dx_prime = torch.split(dy_dx, 1, -1)
    
    laplace = [torch.autograd.grad(g.sum(), xt, retain_graph=True)[0] for g, xt in zip(dy_dx_prime, x_flat)]
    
    return torch.stack(laplace)

def function(x):
    w = torch.normal(0., 1., (*x.shape[1:][::-1],), device=x.device)
    return (x @ w).tanh()
   
m, n, p = 100, 4, 10
x = torch.normal(0., 1., (m, n, p)).to('cuda')
grads(x)  # returns result

```
<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-19-aff8ec7b5d97> in <module>
     11 wf = fermiNet(mol, n_det=4, n_sh=64, n_ph=16, diagonal=False)
     12 pretrainer = Pretrainer(mol)
---> 13 pretrainer.pretrain(wf, n_it=100000)
     14 

~/projects/nn_ansatz/src/pytorch/pretraining_v2.py in pretrain(self, wf, n_walkers, n_it, lr)
    133 
    134         for step in steps:
--> 135             e_locs = compute_local_energy(wf, walkers, self.mol.r_atoms,  self.mol.z_atoms)
    136 
    137             walkers = sampler(wf, self, walkers)

~/projects/nn_ansatz/src/pytorch/vmc.py in compute_local_energy(wf, walkers, r_atoms, z_atoms)
     50 def compute_local_energy(wf, walkers, r_atoms, z_atoms):
     51 
---> 52     first_order_squared, second_order = laplacian(wf, walkers)
     53     potential_energy = compute_potential_energy(walkers, r_atoms, z_atoms)
     54 

~/projects/nn_ansatz/src/pytorch/vmc.py in laplacian(model, walkers)
     69     grads = tc.split(dlogphi_dr, 1, -1)
     70 
---> 71     d2logphi_dr2 = tc.stack([tc.autograd.grad(grad.sum(), r, retain_graph=True)[0] for grad, r in zip(grads, r_s)], 1)
     72 
     73     return dlogphi_dr**2, d2logphi_dr2

~/projects/nn_ansatz/src/pytorch/vmc.py in <listcomp>(.0)
     69     grads = tc.split(dlogphi_dr, 1, -1)
     70 
---> 71     d2logphi_dr2 = tc.stack([tc.autograd.grad(grad.sum(), r, retain_graph=True)[0] for grad, r in zip(grads, r_s)], 1)
     72 
     73     return dlogphi_dr**2, d2logphi_dr2

~/.local/lib/python3.8/site-packages/torch/autograd/__init__.py in grad(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused)
    221         retain_graph = create_graph
    222 
--> 223     return Variable._execution_engine.run_backward(
    224         outputs, grad_outputs_, retain_graph, create_graph,
    225         inputs, allow_unused, accumulate_grad=False)

RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasSgemmStridedBatched( handle, opa, opb, m, n, k, &alpha, a, lda, stridea, b, ldb, strideb, &beta, c, ldc, stridec, num_batches)`

```
## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
Does not fail

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)

Collecting environment information...
PyTorch version: 1.8.0+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 11.1.74
GPU models and configuration: GPU 0: GeForce GTX 1070 with Max-Q Design
Nvidia driver version: 455.23.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.1
[pip3] torch==1.8.0+cu111
[pip3] torchaudio==0.8.0
[pip3] torchvision==0.9.0+cu111
[conda] Could not collect


## Additional context

<!-- Add any other context about the problem here. -->

```
(goku) xmax@deus-ex:/usr/include$ cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2
#define CUDNN_MAJOR 8
#define CUDNN_MINOR 1
#define CUDNN_PATCHLEVEL 1
--
#define CUDNN_VERSION (CUDNN_MAJOR * 1000 + CUDNN_MINOR * 100 + CUDNN_PATCHLEVEL)

#endif /* CUDNN_VERSION_H */
(goku) xmax@deus-ex:/usr/include$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2020 NVIDIA Corporation
Built on Tue_Sep_15_19:10:02_PDT_2020
Cuda compilation tools, release 11.1, V11.1.74
Build cuda_11.1.TC455_06.29069683_0
(goku) xmax@deus-ex:/usr/include$ 
```




cc @ngimel