Rishi Puri

okay id assume @dszwicht will PR that to pyg_lib. in that case then `self._WITH_PYG_LIB = torch.cuda.is_available() and _WITH_PYG_LIB ` should change to `self._WITH_PYG_LIB = _WITH_PYG_LIB` then once pyg_lib is updated...

https://github.com/pyg-team/pyg-lib/pull/111 this should be merged now since the associated PR from intel is in

**Candidates for first round of integration:** - [RGCNConv](https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/conv/rgcn_conv.html#RGCNConv ) - [HeteroLinear](https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/dense/linear.html#HeteroLinear ) **Candidates for second round(`nvfuser+pyg-lib`):** - [HeteroConv](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html?highlight=heteroconv#torch_geometric.nn.conv.HeteroConv)

> Important to note, this will require adding backward support of segment_matmul in pyg-lib. https://github.com/pyg-team/pyg-lib/pull/61 https://github.com/pyg-team/pyg-lib/pull/63

https://github.com/pyg-team/pytorch_geometric/pull/5052

speeding it up: https://github.com/pyg-team/pyg-lib/pull/73

heterolinear cutlass support: https://github.com/pyg-team/pytorch_geometric/pull/5330

@Kh4L reviews welcome

@akihironitta @rusty1s addressed reviews, plz lmk if anything else needed to merge.

@yanbing-j if not opposed I can take this over when I find time in the next few weeks and finish this PR as it is needed for my work