A potencial solution to enable hardware acceleration on torchdrug

[ccorbi]

Hello, I recently installed torch drug on my Apple Laptop with Apple silicon. And I realize that Torchdrug was not taking advantage of the hardware acceleration. I manage to do it so by modifying torch drug , see below the details. It improves significantly the performance of torch drug. I am wondering if you would be interested on adding this modification to the code base?

on line 83 at torchdrug.core.engine.py

        if gpus is None:
            if torch.backends.mps.is_available():
                self.device = torch.device("mps")  
            else:
                self.device = torch.device("cpu")  
        else:
            if len(gpus) != self.world_size:

I am aware that this lines are not compatible with torch <1.13, so may be it will work better to encapsule that in try/except . Also not sure if I should include this in other parts of the code. What do you think about it ?