[PERFORMANCE] Use index_select instead of regular indexing

Hi all, while profiling my code, I found that using regular indexing makes the backward pass >100x slower than the forward pass. Apparently doing something like x[indices] is very slow if there are many repeated indices. However, using x.index_select(dim, indices) fixes the issue. This seems to be very relevant for GNNs, since there are many node/edges/emb indexing.

I am not sure this is fixed in newer versions of pytorch or if this is standard practice in DGL implementation, but thought of sharing, since I would never expect pytorch to have such characteristic.

Reference: https://github.com/pytorch/pytorch/issues/41162