row.device().is_cpu() INTERNAL ASSERT FAILED at "csrc/cpu/diag_cpu.cpp":7
Labels: triaged, module: macos

### 🐛 Describe the bug

[gist](https://gist.github.com/inoue0426/36baa6ec2f87d537e16b1029080654ed)

I got the runtime error related to the M1 Mac, although this code works for Linux with GTX3080.
(This may be related to PyTorch Geometric, not PyTorch.)

```python
/Users/yoshitakainoue/.pyenv/versions/miniforge3-4.10.1-5/lib/python3.9/site-packages/torch_sparse/utils.py:20: UserWarning: MPS: no support for int64 min/max ops, casting it to int32 (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/operations/Sort.mm:39.)
  return inputs.sort()
/Users/yoshitakainoue/.pyenv/versions/miniforge3-4.10.1-5/lib/python3.9/site-packages/torch_sparse/storage.py:69: UserWarning: MPS: no support for int64 min/max ops, casting it to int32 (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/operations/ReduceOps.mm:1271.)
  assert trust_data or int(row.max()) < M
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[16], line 16
     13 model.train()
     14 optimizer.zero_grad()
---> 16 outputs, all_attention, result = model(x, adj, train_drug, train_cell)
     17 torch.cuda.empty_cache()
     18 loss = criterion(outputs.squeeze(), train_labels)

File ~/.pyenv/versions/miniforge3-4.10.1-5/lib/python3.9/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

Cell In[12], line 17, in GAT.forward(self, x, adj, drug, cell)
     16 def forward(self, x, adj, drug, cell):
---> 17     x, attention = self.gat1(x, adj.t(), return_attention_weights=True)
     18     all_attention = get_attention_mat(attention)
     19     del attention 

File ~/.pyenv/versions/miniforge3-4.10.1-5/lib/python3.9/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File ~/.pyenv/versions/miniforge3-4.10.1-5/lib/python3.9/site-packages/torch_geometric/nn/conv/gat_conv.py:244, in GATConv.forward(self, x, edge_index, edge_attr, size, return_attention_weights)
    242 elif isinstance(edge_index, SparseTensor):
    243     if self.edge_dim is None:
--> 244         edge_index = torch_sparse.set_diag(edge_index)
    245     else:
    246         raise NotImplementedError(
    247             "The usage of 'edge_attr' and 'add_self_loops' "
    248             "simultaneously is currently not yet supported for "
    249             "'edge_index' in a 'SparseTensor' form")

File ~/.pyenv/versions/miniforge3-4.10.1-5/lib/python3.9/site-packages/torch_sparse/diag.py:41, in set_diag(src, values, k)
     38 src = remove_diag(src, k=k)
     39 row, col, value = src.coo()
---> 41 mask = torch.ops.torch_sparse.non_diag_mask(row, col, src.size(0),
     42                                             src.size(1), k)
     43 inv_mask = ~mask
     45 start, num_diag = -k if k < 0 else 0, mask.numel() - row.numel()

File ~/.pyenv/versions/miniforge3-4.10.1-5/lib/python3.9/site-packages/torch/_ops.py:502, in OpOverloadPacket.__call__(self, *args, **kwargs)
    497 def __call__(self, *args, **kwargs):
    498     # overloading __call__ to ensure torch.ops.foo.bar()
    499     # is still callable from JIT
    500     # We save the function ptr as the `op` attribute on
    501     # OpOverloadPacket to access it here.
--> 502     return self._op(*args, **kwargs or {})

RuntimeError: row.device().is_cpu() INTERNAL ASSERT FAILED at "csrc/cpu/diag_cpu.cpp":7, please report a bug to PyTorch. row must be CPU tensor
```

### Versions

```
Collecting environment information...
PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.2.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: version 3.23.3
Libc version: N/A

Python version: 3.9.5 | packaged by conda-forge | (default, Jun 19 2021, 00:24:55)  [Clang 11.1.0 ] (64-bit runtime)
Python platform: macOS-13.2.1-arm64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1

Versions of relevant libraries:
[pip3] flake8==5.0.3
[pip3] numpy==1.23.1
[pip3] torch==2.0.1
[pip3] torch-geometric==2.3.1
[pip3] torch-scatter==2.1.1
[pip3] torch-sparse==0.6.17
[pip3] torchaudio==2.0.2
[pip3] torchvision==0.15.2
[conda] numpy                     1.23.1                   pypi_0    pypi
[conda] torch                     2.0.1                    pypi_0    pypi
[conda] torch-geometric           2.3.1                    pypi_0    pypi
[conda] torch-scatter             2.1.1                    pypi_0    pypi
[conda] torch-sparse              0.6.17                   pypi_0    pypi
[conda] torchaudio                2.0.2                    pypi_0    pypi
[conda] torchvision               0.15.2                   pypi_0    pypi
```

cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr @abhudev