Repository: pyg-team/pytorch_geometric

Issue:
If RGATConv uses f-scaled mod, there will be bugs of different tensors on different devices
### 🐛 Describe the bug

When using RGATConv, set mod = 'f-scaled', tensor in different devices error will appear .


`RGATConv(in_channels, out_channels, 14, num_bases=8,mod='f-scaled',concat =False,heads=8)` 

`RGATConv(in_channels, out_channels, 14, num_bases=8,concat =False,heads=8,)`,


error:

>   File "/home/mist/tianchi/RevRgcn.py", line 365, in forward
    return self.conv(x, edge_index,edge_type)
  File "/mistgpu/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/mistgpu/site-packages/torch_geometric/nn/conv/rgat_conv.py", line 341, in forward
    out = self.propagate(edge_index=edge_index, edge_type=edge_type, x=x,
  File "/mistgpu/site-packages/torch_geometric/nn/conv/message_passing.py", line 374, in propagate
    out = self.message(**msg_kwargs)
  File "/mistgpu/site-packages/torch_geometric/nn/conv/rgat_conv.py", line 473, in message
    degree = scatter_add(ones, index,
  File "/mistgpu/site-packages/torch_scatter/scatter.py", line 29, in scatter_add
    return scatter_sum(src, index, dim, out, dim_size)
  File "/mistgpu/site-packages/torch_scatter/scatter.py", line 21, in scatter_sum
    return out.scatter_add_(dim, index, src)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper_scatter_add_)


### Environment

* PyG version:`2.0.4`
* PyTorch version:`1.12.1`
* OS:`ubuntu 18.04.6 LTS`
* Python version:`3.9.12`
* CUDA/cuDNN version:`cu116`
* How you installed PyTorch and PyG (`conda`, `pip`, source): `pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-1.12.0+cu116.html`
* Any other relevant information (*e.g.*, version of `torch-scatter`):



Discussion:
Thanks for reporting. Will take a look!

Resolving patch:
diff --git a/torch_geometric/nn/conv/rgat_conv.py b/torch_geometric/nn/conv/rgat_conv.py
--- a/torch_geometric/nn/conv/rgat_conv.py
+++ b/torch_geometric/nn/conv/rgat_conv.py
@@ -441,7 +441,7 @@ def message(self, x_i: Tensor, x_j: Tensor, edge_type: Tensor,
 
         elif self.mod == "scaled":
             if self.attention_mode == "additive-self-attention":
-                ones = torch.ones(index.size())
+                ones = alpha.new_ones(index.size())
                 degree = scatter_add(ones, index,
                                      dim_size=size_i)[index].unsqueeze(-1)
                 degree = torch.matmul(degree, self.l1) + self.b1
@@ -453,7 +453,7 @@ def message(self, x_i: Tensor, x_j: Tensor, edge_type: Tensor,
                     alpha.view(-1, self.heads, 1),
                     degree.view(-1, 1, self.out_channels))
             elif self.attention_mode == "multiplicative-self-attention":
-                ones = torch.ones(index.size())
+                ones = alpha.new_ones(index.size())
                 degree = scatter_add(ones, index,
                                      dim_size=size_i)[index].unsqueeze(-1)
                 degree = torch.matmul(degree, self.l1) + self.b1
@@ -469,7 +469,7 @@ def message(self, x_i: Tensor, x_j: Tensor, edge_type: Tensor,
             alpha = torch.where(alpha > 0, alpha + 1, alpha)
 
         elif self.mod == "f-scaled":
-            ones = torch.ones(index.size())
+            ones = alpha.new_ones(index.size())
             degree = scatter_add(ones, index,
                                  dim_size=size_i)[index].unsqueeze(-1)
             alpha = alpha * degree


Test patch: