Support Running GNNs on Specific GPU

Hi DeepRobust Team,

I’ve encountered a problem when testing GNNs on devices other than cuda:0. To reproduce, simply modifying line 18 in test_gcn.py to

device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")

Then the error message is prompted as follows:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices

I found it’s caused by normalize_adj_tensor and degree_normalize_adj_tensor, where the device is set to be cuda:0 by default wherever adj is. Though many people run experiments on cuda:0 by default, it could be even better to support running GNNs in DeepRobust on other devices specified via the device variable.

Could you add this feature in the next update? Thank you. 😀

[ChandlerBang]

Sorry for the late reply. I have just fixed the issue. Thank you @LFhase and @EdisonLeeeee !

[LFhase]

Thank you @ChandlerBang and @EdisonLeeeee. I think @EdisonLeeeee 's solution would do the job if adj.device can be always accessed in each call of the two functions. Would you like to consider updating the two lines if it’s feasible? 😃

BTW, I don’t have further questions.