Repository: pytorch
Issue: CUDA error: invalid configuration argument for torch.sparse tensor backward
Labels: module: sparse,module: cuda,triaged

## 🐛 Bug

When using backward for sparse tensor on CUDA, which works well for both dense+CUDA and sparse+CPU, an CUDA error occurs: invalid configuration argument .

## To Reproduce

Steps to reproduce the behavior:

1. sparse+CUDA(error occurs)
```python
>>> import torch
>>> torch.cuda.set_device(1)
>>> S = torch.sparse.FloatTensor(torch.LongTensor([[1],[0]]), torch.ones(1), [2, 2])
>>> d = torch.nn.Parameter(S.cuda())
>>> loss = torch.sparse.sum(d**2)
>>> loss.backward()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/chenzhb/anaconda3/envs/torch/lib/python3.8/site-packages/torch/_tensor.py", line 255, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/home/chenzhb/anaconda3/envs/torch/lib/python3.8/site-packages/torch/autograd/__init__.py", line 147, in backward
    Variable._execution_engine.run_backward(
RuntimeError: CUDA error: invalid configuration argument
>>> d.grad
>>> 
```
2. dense+CUDA(fine)
```python
>>> a = torch.nn.Parameter(torch.rand(3,3).cuda())
>>> loss = (a**2).sum()
>>> loss.backward()
>>> a.grad
tensor([[0.3801, 1.6425, 1.1757],
        [0.8900, 1.1500, 1.6633],
        [0.3145, 0.7741, 1.7915]], device='cuda:1')
```
3. sparse+CPU(fine)
```python
>>> dd = torch.nn.Parameter(S)
>>> loss = torch.sparse.sum(dd**2)
>>> loss.backward()
>>> dd.grad
tensor(indices=tensor([[1],
                       [0]]),
       values=tensor([2.]),
       size=(2, 2), nnz=1, layout=torch.sparse_coo)
```

## Expected behavior

There should not be error in backward, and the grad should be calculated correctly.

## Environment
PyTorch version: 1.9.0+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.3 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.27

Python version: 3.8.10 (default, Jun  4 2021, 15:09:15)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.15.0-55-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: A100-PCIE-40GB
GPU 1: A100-PCIE-40GB
GPU 2: A100-PCIE-40GB

Nvidia driver version: 460.80
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.3
[pip3] pytorch-transformers==1.2.0
[pip3] torch==1.9.0+cu111
[pip3] torch-cluster==1.5.9
[pip3] torch-geometric==1.7.2
[pip3] torch-scatter==2.0.6
[pip3] torch-sparse==0.6.9
[pip3] torch-spline-conv==1.2.1
[pip3] torchaudio==0.9.0
[pip3] torchvision==0.10.0+cu111
[conda] blas                      1.0                         mkl    defaults
[conda] cudatoolkit               11.1.1               h6406543_8    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.2.0           h06a4308_296    defaults
[conda] mkl-service               2.4.0            py38h497a2fe_0    conda-forge
[conda] mkl_fft                   1.3.0            py38h42c9631_2    defaults
[conda] mkl_random                1.2.2            py38h1abd341_0    conda-forge
[conda] numpy                     1.20.3                   pypi_0    pypi
[conda] numpy-base                1.20.2           py38hfae3a4d_0    defaults
[conda] pytorch-transformers      1.2.0                    pypi_0    pypi
[conda] torch                     1.9.0+cu111              pypi_0    pypi
[conda] torch-cluster             1.5.9                    pypi_0    pypi
[conda] torch-geometric           1.7.2                    pypi_0    pypi
[conda] torch-scatter             2.0.6                    pypi_0    pypi
[conda] torch-sparse              0.6.9                    pypi_0    pypi
[conda] torch-spline-conv         1.2.1                    pypi_0    pypi
[conda] torchaudio                0.9.0                    pypi_0    pypi
[conda] torchvision               0.10.0+cu111             pypi_0    pypi

cc @nikitaved @pearu @cpuhrsch @IvanYashchuk @ngimel