Minor fix for pyg-lib usage in HeteroLinear and RGCNConv

[puririshi98]

pyg_lib pathway is chosen if cuda is available and pyg_lib is available. but if the input is not cuda we should not be using the pyg_lib pathway. simple fix

without this a ton of CI fails for testing rgcnconv and heterolinear w/ cpu inputs

errors.txt

[codecov]

Codecov Report

Merging #5510 (38f2a49) into master (3733006) will increase coverage by 0.02%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master    #5510      +/-   ##
==========================================
+ Coverage   83.84%   83.86%   +0.02%     
==========================================
  Files         349      349              
  Lines       19231    19231              
==========================================
+ Hits        16124    16128       +4     
+ Misses       3107     3103       -4     

Impacted Files	Coverage Δ
torch_geometric/nn/conv/rgcn_conv.py	95.39% <100.00%> (+1.31%)	⬆️
torch_geometric/nn/dense/linear.py	87.21% <100.00%> (+3.75%)	⬆️
torch_geometric/utils/scatter.py	66.66% <0.00%> (-33.34%)	⬇️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[EdisonLeeeee]

Use x.is_cuda for device check?

[rusty1s]

This PR will currently break TorchScript functionality since we cannot check for device within a TorchScript program (that's why we have this awkward

self._WITH_PYG_LIB = torch.cuda.is_available() and _WITH_PYG_LIB

check as a workaround in __init__. I suggest that we simply integrate CPU-based segment_matmul support in pyg-lib via the classic for-loop instead.

[rusty1s]

cc @dszwicht

[DamianSzwichtenberg]

@rusty1s I have had a minimal implementation of [segment | grouped]_matmul ready for some time. Was working on optimization, but in such a case I'll create a PR with basic solution later today 😉

[puririshi98]

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 all pathways should work

[rusty1s]

Indeed. PR to pyg-lib just landed, so this should be resolved soon.

[puririshi98]

pyg-team/pyg-lib#111
this should be merged now since the associated PR from intel is in

[rusty1s]

Thanks for the update @puririshi98. I noticed that segment_matmul still does not support TorchScript because the Python implementation via torch.autograd.Function does not come with TorchScript support. Now that the dispatching of segment_matmul is resolved in pyg-lib, can we move its backward implementation from Python to C++? Let me know if you have time working on this!