Add a non-blocking option on DGLGraph.to

[ZhaofengWu]

🚀 Feature

Change DGLGraphs's to from def to(self, ctx): to def to(self, ctx, non_blocking=False): or with some other default.

Motivation

I'm using pytorch-lightning, a framework that provides a scaffold for many PyTorch applications, with DGL. One thing they do is that they automatically call .to on every parameter to forward (https://github.com/PyTorchLightning/pytorch-lightning/blob/98f7842970fd593c21bebfdba5da8aa311d50730/pytorch_lightning/trainer/distrib_parts.py#L458). They recently added non_blocking=True to this .to call (Lightning-AI/pytorch-lightning#1843), making it no longer compatible with DGLGraph as its .to doesn't have such an argument. Therefore, is it possible that a non_blocking option be supported in DGL too?

Alternatives

Something like def to(self, ctx, **kwargs): with **kwargs entirely ignored. That'll solve the compatibility issue but is not very elegant.

[yzh119]

@ZhaofengWu Thanks for your suggestion, I'll take a look :)

[sophiakrix]

Hi there!

Are you using the dataloaders from DGL with pytorch lightning? I am trying to do that but I get the suspicion that it does not accept the NodeDataLoader and EdgeDataLoaders.

What is your experience?

[BarclayII]

@sophiakrix Currently we are working with PyTorch Lightning to make DGL compatible. It requires some changes in DGL's NodeDataLoader and EdgeDataLoader.

[github-actions]

This issue has been automatically marked as stale due to lack of activity. It will be closed if no further activity occurs. Thank you

[jermainewang]

The issue has been resolved in the latest release. Both Torch Lightning and async copy have been supported.