CuGraphRGCNConv static graph support

[andigu]

The cheatsheet says CuGraphRGCNConv supports static graphs. If I run something like

rgcn = CuGraphRGCNConv(embedding_dim, embedding_dim, num_relations=12).to(device)
X = torch.randn(batch_size, num_nodes, embedding_dim, device=device, dtype=torch.float32)
edges = ... # (number of edges in a single graph, 2)
edge_types = ... # (number of edges in single graph, )
rgcn_model(X, EdgeIndex(edges), edge_types)

throws an error:

  File "timing_comparison.py", line 99, in time_model
    _ = model(*inputs)
        ^^^^^^^^^^^^^^
  File ".venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.12/site-packages/torch_geometric/nn/conv/cugraph/rgcn_conv.py", line 99, in forward
    out = RGCNConvAgg(x, self.comp, graph, concat_own=self.root_weight,
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.12/site-packages/pylibcugraphops/pytorch/operators/agg_hg_basis.py", line 75, in agg_hg_basis_n2n_post
    return _agg_hg_basis_n2n_post_autograd.apply(feat, weights_comb, graph, options)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.12/site-packages/torch/autograd/function.py", line 575, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.12/site-packages/pylibcugraphops/pytorch/operators/agg_hg_basis.py", line 102, in forward
    agg_hg_basis_n2n_post_fwd(
ValueError: input_embedding expected 2 dimensions but got 3

Is this the wrong way to use RGCN with static graphs?