Fix circular import in GPSE (#10190)

Co-authored-by: Rishi Puri <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

torch_geometric/nn/models/gpse.py

@@ -9,6 +9,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch.nn import Module
 from tqdm import trange
 
 import torch_geometric.transforms as T
@@ -715,8 +716,9 @@ def forward(self, x, pos_enc):
 
 
 @torch.no_grad()
-def gpse_process(model: GPSE, data: Data, rand_type: str, use_vn: bool = True,
-                 bernoulli_thresh: float = 0.5, neighbor_loader: bool = False,
+def gpse_process(model: Module, data: Data, rand_type: str,
+                 use_vn: bool = True, bernoulli_thresh: float = 0.5,
+                 neighbor_loader: bool = False,
                  num_neighbors: List[int] = [30, 20, 10], fillval: int = 5,
                  layers_mp: int = None, **kwargs) -> torch.Tensor:
     r"""Processes the data using the :class:`GPSE` model to generate and append
@@ -731,7 +733,7 @@ def gpse_process(model: GPSE, data: Data, rand_type: str, use_vn: bool = True,
     :obj:`precompute_GPSE` on your whole dataset is advised instead.
 
     Args:
-        model (GPSE): The :class:`GPSE` model.
+        model (Module): The :class:`GPSE` model.
         data (torch_geometric.data.Data): A :class:`~torch_geometric.data.Data`
             object.
         rand_type (str, optional): Type of random features to use. Options are

torch_geometric/transforms/add_gpse.py

@@ -1,6 +1,7 @@
+from torch.nn import Module
+
 from torch_geometric.data import Data
 from torch_geometric.data.datapipes import functional_transform
-from torch_geometric.nn.models.gpse import GPSE
 from torch_geometric.transforms import BaseTransform, VirtualNode
 
 
@@ -13,15 +14,15 @@ class AddGPSE(BaseTransform):
     the actual encodings.
 
     Args:
-        model (GPSE): The pre-trained GPSE model.
+        model (Module): The pre-trained GPSE model.
         use_vn (bool, optional): Whether to use virtual nodes.
             (default: :obj:`True`)
         rand_type (str, optional): Type of random features to use. Options are
             :obj:`NormalSE`, :obj:`UniformSE`, :obj:`BernoulliSE`.
             (default: :obj:`NormalSE`)
 
     """
-    def __init__(self, model: GPSE, use_vn: bool = True,
+    def __init__(self, model: Module, use_vn: bool = True,
                  rand_type: str = 'NormalSE'):
         self.model = model
         self.use_vn = use_vn