Eliminate warning responsible for nondeterministic test failure

CHANGELOG.md

@@ -21,6 +21,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Fixed
 
+- Fixed `test_to_hetero_lazy_cuda` by removing nondeterministi behaviour ([#10531](https://github.com/pyg-team/pytorch_geometric/pull/10531))
 - Fixed `ogbn_train_cugraph` example for distributed cuGraph ([#10439](https://github.com/pyg-team/pytorch_geometric/pull/10439))
 - Added `safe_onnx_export` function with workarounds for `onnx_ir.serde.SerdeError` issues in ONNX export ([#10422](https://github.com/pyg-team/pytorch_geometric/pull/10422))
 - Fixed importing PyTorch Lightning in `torch_geometric.graphgym` and `torch_geometric.data.lightning` when using `lightning` instead of `pytorch-lightning` ([#10404](https://github.com/pyg-team/pytorch_geometric/pull/10404), [#10417](https://github.com/pyg-team/pytorch_geometric/pull/10417)))

test/nn/test_to_hetero_transformer.py

@@ -557,6 +557,13 @@ def forward(self, x, edge_index, edge_attr):
             return self.conv(x, edge_index, edge_attr)
 
     data = FakeHeteroDataset(edge_dim=10)[0].to('cuda')
+
+    # Ensure each node type is a destination at least once
+    for node_type in data.node_types:
+        idx = torch.arange(data[node_type].num_nodes, device='cuda')
+        edge_index = torch.stack([idx, idx])
+        data[(node_type, 'self', node_type)].edge_index = edge_index
+
     model = to_hetero(Model(), data.metadata())
     out_dict = model(data.x_dict, data.edge_index_dict, data.edge_attr_dict)
     assert len(out_dict) == len(data.node_types)

torch_geometric/nn/pool/connect/base.py

@@ -48,9 +48,6 @@ def __init__(
         self.batch = batch
 
 
-ConnectOutput = torch.jit.script(ConnectOutput)
-
-
 class Connect(torch.nn.Module):
     r"""An abstract base class for implementing custom edge connection
     operators as described in the `"Understanding Pooling in Graph Neural

torch_geometric/nn/pool/select/base.py

@@ -62,9 +62,6 @@ def __init__(
         self.weight = weight
 
 
-SelectOutput = torch.jit.script(SelectOutput)
-
-
 class Select(torch.nn.Module):
     r"""An abstract base class for implementing custom node selections as
     described in the `"Understanding Pooling in Graph Neural Networks"