When using GAT, how could I update a subset of nodes, instead of all nodes?

[RenqinCai]

I am using GAT for a node classification task. In my task, there are multiple different types of nodes. I would like to update one type of node after a message passing, instead of updating all nodes simultaneously. Any thoughts? Thanks.

[wsad1]

A simple thing to do would be, to run message passing on all nodes and then only update nodes of one type.

x=... # node feature matrix.
node_type=... # 1-d tensor indicating the type of each node.
node_mask = node_type ==i
conv =GATConv()
x[node_mask] = conv(x,edge_index)[node_mask] # update nodes of type i.

[RenqinCai]

Thanks for your reply. I would give it a try.