Repository: pytorch
Issue: torch.gather with sparse_grad=True does not work with SGD optimizer with momentum; gives bad error message
Labels: module: sparse,module: optimizer,module: error checking,triaged,actionable

## 🐛 Bug

`torch.gather` with `sparse_grad=True` does not work with SGD optimizer with momentum.

The following example does output the batch 0 data, but cannot run when batch id > 1.

## To Reproduce

```Python
import torch 
import torch.nn as nn
import torch.nn.functional as F

SPARSE_GRAD = True

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.weight = nn.Parameter(torch.rand(2, 10)) 

    def forward(self, x):
        index = torch.tensor([[7, 4, 5], [0, 1, 2]])
        w = torch.gather(self.weight, dim=1, index=index, sparse_grad=SPARSE_GRAD) 
        loss = F.linear(x, w)
        return loss.mean()

net = Net()
params = net.parameters()
optimizer = torch.optim.SGD(params, lr=1e-3, momentum=0.9) 

for i in range(10):
    x = torch.rand(32, 3) 
    loss = net(x)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    print('Batch {}, Loss {}'.format(i, loss.detach().cpu().numpy()))
```

Error output:
```
Batch 0, Loss 1.0152490139007568
Traceback (most recent call last):
  File "torch_gather_test.py", line 26, in <module>
    optimizer.step()
  File "/work/.torch-1.9.0/lib/python3.8/site-packages/torch/optim/optimizer.py", line 88, in wrapper
    return func(*args, **kwargs)
  File "/work/.torch-1.9.0/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "/work/.torch-1.9.0/lib/python3.8/site-packages/torch/optim/sgd.py", line 110, in step
    F.sgd(params_with_grad,
  File "/work/.torch-1.9.0/lib/python3.8/site-packages/torch/optim/_functional.py", line 173, in sgd
    buf.mul_(momentum).add_(d_p, alpha=1 - dampening)
RuntimeError: set_indices_and_values_unsafe is not allowed on a Tensor created from .data or .detach().
If your intent is to change the metadata of a Tensor (such as sizes / strides / storage / storage_offset)
without autograd tracking the change, remove the .data / .detach() call and wrap the change in a `with torch.no_grad():` block.
For example, change:
    x.data.set_(y)
to:
    with torch.no_grad():
        x.set_(y)
```


## Expected behavior

Should work both when `SPARSE_GRAD=True` and `SPARSE_GRAD=False`.

## Environment

PyTorch version: 1.9.0+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (GCC) 8.2.0
Clang version: 3.8.0 (tags/RELEASE_380/final)
CMake version: version 3.16.0
Libc version: glibc-2.27

Python version: 3.8.8 (default, Feb 20 2021, 21:09:14)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-3.10.0_3-0-0-17-x86_64-with-glibc2.27
Is CUDA available: False
CUDA runtime version: 11.2.142
GPU models and configuration: 
GPU 0: Tesla V100-PCIE-16GB
GPU 1: Tesla V100-PCIE-16GB
GPU 2: Tesla V100-PCIE-16GB
GPU 3: Tesla V100-PCIE-16GB
GPU 4: Tesla V100-PCIE-16GB
GPU 5: Tesla V100-PCIE-16GB
GPU 6: Tesla V100-PCIE-16GB
GPU 7: Tesla V100-PCIE-16GB

Nvidia driver version: 450.51.06
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.1.1
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.21.1
[pip3] torch==1.9.0+cpu
[pip3] torchaudio==0.9.0
[pip3] torchvision==0.10.0+cpu
[conda] Could not collect

## Additional context

<!-- Add any other context about the problem here. -->


cc @nikitaved @pearu @cpuhrsch @amjames @bhosmer @vincentqb @jbschlosser @albanD @IvanYashchuk @iramazanli