For HeteroData, can we assign edge_type to only one set of the edges?

[ZihanChen1995]

Suppose we have a HeteroData like following. Our target is using GNN to predict the links between A and B (i.e., (A, a_to_b, B)). I have already built a working model using to_hetero():

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.conv1 = SAGEConv(768, hidden_channels)
        self.conv2 = SAGEConv(768, hidden_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

gnn = GNN(hidden_channels)
gnn = to_hetero(self.gnn, metadata=all_data.metadata())

This model works well. However, since the edge (A, a_to_b, B) has its unique edge_type (3 relations as 0,1,2), I am wondering:

Can we utilize this edge_type information?

I tried to follow the example here by assigning RGCNConv to the (A, a_to_b, B).

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, metadata, num_layers):
        super().__init__()

        self.convs = torch.nn.ModuleList()
        
        for _ in range(num_layers): 
            
            temd = {} # temp dict to store layers

            for edge_type in metadata[1]:

                if edge_type==('A', 'a_to_b', 'B'):
                    temd[edge_type] = RGCNConv(in_channels=768, 
                                               out_channels=hidden_channels, 
                                               num_relations=3)
                else:
                    temd[edge_type] = SAGEConv(in_channels=768, 
                                               out_channels=hidden_channels, 
                                               )
            conv = HeteroConv(temd)
            self.convs.append(conv)

    def forward(self, x_dict, edge_index_dict, edge_type_dict):
        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict, edge_type_dict)
            x_dict = {key: F.leaky_relu(x) for key, x in x_dict.items()}
        return x_dict

model = GNN(hidden_channels=64, metadata=all_data.metadata(), num_layers=2)
model = model.to(device)
model = model.double()

However, it keeps showing the error

[/usr/local/lib/python3.9/dist-packages/torch/_ops.py](https://localhost:8080/#) in __call__(self, *args, **kwargs)
    440         # We save the function ptr as the `op` attribute on
    441         # OpOverloadPacket to access it here.
--> 442         return self._op(*args, **kwargs or {})
    443 
    444     # TODO: use this to make a __dir__

RuntimeError: expected scalar type Float but found Double

Would you mind giving me some suggestions about how I can utilize the edge_type information of 'a_to_b'? I guess a naive way is splitting 'a_to_b' into 'a_to_b_rel1', 'a_to_b_rel2', 'a_to_b_rel3'. But do you think is there any better solution?

Thank you so much in advance!

---

Here is the HereoData:

HeteroData(
  A={
    node_id=[50000],
    x=[50000, 768]
  },
  B={
    node_id=[4000],
    x=[4000, 768]
  },
  (A, a_to_a, A)={ edge_index=[2, 96532] },
  (B, b_to_b, B)={ edge_index=[2, 39634] },
  (A, a_to_b, B)={
    edge_index=[2, 30000],
    edge_type=[30000],
    edge_label=[10000],
    edge_label_index=[2, 10000]
  }

[rusty1s]

I think leveraging RGCNConv for this single edge type is reasonable. Do you have a small reproducible example? Are there any values of dtype==torch.double in your data?