[GraphBolt][CUDA] Link prediction `--mode=cuda-cuda` fix.

## Description Peak memory usage after fixing bug is 22GB (was getting runtime error before). Any GPU with 24GB should be able to run the `cuda-cuda` mode successfully. Error message: ``` Traceback (most recent call last): File "/home/mfbalin/dgl-1/examples/sampling/graphbolt/lightning/../link_prediction.py", line 427, in <module> main(args) File "/home/mfbalin/dgl-1/examples/sampling/graphbolt/lightning/../link_prediction.py", line 416, in main valid_mrr, test_mrr = evaluate( File "/home/mfbalin/.local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context return func(*args, **kwargs) File "/home/mfbalin/dgl-1/examples/sampling/graphbolt/lightning/../link_prediction.py", line 297, in evaluate compute_mrr(args, model, node_emb, seeds, labels, indexes) File "/home/mfbalin/.local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context return func(*args, **kwargs) File "/home/mfbalin/dgl-1/examples/sampling/graphbolt/lightning/../link_prediction.py", line 272, in compute_mrr return mrr(preds, labels, indexes=indexes) File "/home/mfbalin/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl return self._call_impl(*args, **kwargs) File "/home/mfbalin/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl return forward_call(*args, **kwargs) File "/home/mfbalin/.local/lib/python3.10/site-packages/torchmetrics/metric.py", line 290, in forward self._forward_cache = self._forward_reduce_state_update(*args, **kwargs) File "/home/mfbalin/.local/lib/python3.10/site-packages/torchmetrics/metric.py", line 358, in _forward_reduce_state_update batch_val = self.compute() File "/home/mfbalin/.local/lib/python3.10/site-packages/torchmetrics/metric.py", line 593, in wrapped_func value = _squeeze_if_scalar(compute(*args, **kwargs)) File "/home/mfbalin/.local/lib/python3.10/site-packages/torchmetrics/retrieval/base.py", line 125, in compute preds = preds[indices] RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu) ``` ## Checklist Please feel free to remove inapplicable items for your PR. - [ ] The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]]) - [ ] I've leverage the [tools](https://docs.google.com/document/d/1iHyj7zlmygKSk5gBPsqIqL5ASPzJSPREaNT_QdsiYA4/edit) to beautify the python and c++ code. - [ ] The PR is complete and small, read the [Google eng practice (CL equals to PR)](https://google.github.io/eng-practices/review/developer/small-cls.html) to understand more about small PR. In DGL, we consider PRs with less than 200 lines of core code change are small (example, test and documentation could be exempted). - [ ] All changes have test coverage - [ ] Code is well-documented - [ ] To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change - [ ] Related issue is referred in this PR - [ ] If the PR is for a new model/paper, I've updated the example index [here](../examples/README.md). ## Changes <!-- You could use following template - [ ] Feature1, tests, (and when applicable, API doc) - [ ] Feature2, tests, (and when applicable, API doc) -->