Support for LightningModules in to_hetero() and to_hetero_with_bases()

[redfave]

🚀 The feature, motivation and pitch

I'm using PyG with PyTorch Lightning.
In PyTorch Lightning the inheritance base is lightning.pytorch.core.LightningModule instead of the plain torch.nn.Module (LightningModule is derived from Module through).

As a consequence, I stumbled upon the problem, that I can't use to_hetero() (also to_hetero_with_bases()) to convert my homogenous GNN into a heterogenous one. The reason is that those methods assume a Module as base:

pytorch_geometric/torch_geometric/nn/to_hetero_transformer.py

Line 137 in a89e022

module: Module,

          model = GNNGraphConv( # own GNN based on LightningModule
              # in_channels=graph.num_node_features,
              in_channels=-1,  # lazy init for heterogenous graphs
              # (channels inferred on first forward pass)
              hidden_channels=trial_params.model_parameters.hidden_channels,
              dropout_rate=trial_params.model_parameters.dropout_rate,
              num_out_classes=graph.num_classes,
              learning_rate=trial_params.learning_parameters.learning_rate,
              weight_decay=trial_params.learning_parameters.weight_decay,
          )
          model = to_hetero(module=model, metadata=graph.metadata(), debug=True)
          trainer.fit(model=model, datamodule=datamodule)

TypeError: `model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `GraphModule.__new__.<locals>.GraphModuleImpl`

A future support for PyTorch Lightning would be very lovely and appreciated.
In the meantime, I think there should be a warning in the documentation, that those methods only support modules directly derived from torch.nn.Module.

Alternatives

No response

Additional context

Lighning issue: Lightning-AI/pytorch-lightning#21629