Build KNN graph over some subset of Node features

I have a point cloud that I want to use a graph neural network on. Each point in the point cloud is characterised by its positional coordinates as well as it’s color. So a single node is (X, Y, Z, C).

Now I want to apply an Edge Convolution on this (as described in the DGL Edge-Conv example, and to do it I should build a Nearest Neighbors graph on (X, Y, Z) (And not on C), then use all the 4 properties as features for my neural network.

What would be a clean and efficient way to do this? (I have a lot of data so I want to batch and collate well)

[This is the cross post of an identical Stack Overflow question:]

Can you use knn_graph?