Repository: THUDM/CogDL (language: python)

Issue:
piplines的gen-emb模型有部分bug
## 🐛 Bug

1.在使用GNN模型的时候如果不传入x的话，piplines里面不会生成随机特征来对齐进行填充
https://github.com/THUDM/cogdl/blob/012669418e60078204f0fb70cb09f614624d99f9/cogdl/pipelines.py#L190-L192
这里生成的随机特征没有赋值给x

2.在跑line模型的时候其train方法出现了循环调用，导致了无限递归
https://github.com/THUDM/cogdl/blob/012669418e60078204f0fb70cb09f614624d99f9/cogdl/models/emb/line.py#L68-L69

3.在跑dgi，grace，unsup_graphsage的时候会出现tensor不在同一个device上的情况

4.在跑node2vec的时候会报错：NotImplementedError: `None` data_wrapper is not implemented.

## To Reproduce
使用官方给的例子


Discussion / hints:
Hi @HaSai666, 非常感谢指出这些bug。我们最近对cogdl整体框架进行了较大幅度的修改，所以引入了一些bug，我们会尽快修复。

Fix patch:
diff --git a/cogdl/data/data.py b/cogdl/data/data.py
index 8acdbb06..9fcf714a 100644
--- a/cogdl/data/data.py
+++ b/cogdl/data/data.py
@@ -676,6 +676,10 @@ def keys(self):
         keys = [key for key in keys if key[:2] != "__" and key[-2:] != "__"]
         return keys
 
+    @property
+    def device(self):
+        return self._adj.device
+
     def degrees(self):
         return self._adj.degrees()
 
diff --git a/cogdl/models/base_model.py b/cogdl/models/base_model.py
index a45293d3..02a688d5 100644
--- a/cogdl/models/base_model.py
+++ b/cogdl/models/base_model.py
@@ -15,7 +15,6 @@ def build_model_from_args(cls, args):
 
     def __init__(self):
         super(BaseModel, self).__init__()
-        self.device = ""
         self.loss_fn = None
         self.evaluator = None
 
@@ -25,8 +24,9 @@ def _forward_unimplemented(self, *input: Any) -> None:  # abc warning
     def forward(self, *args):
         raise NotImplementedError
 
-    def set_device(self, device):
-        self.device = device
+    @property
+    def device(self):
+        return next(self.parameters()).device
 
     def set_loss_fn(self, loss_fn):
         self.loss_fn = loss_fn
diff --git a/cogdl/models/emb/dngr.py b/cogdl/models/emb/dngr.py
index 24d5a8a9..231be939 100644
--- a/cogdl/models/emb/dngr.py
+++ b/cogdl/models/emb/dngr.py
@@ -79,7 +79,7 @@ def __init__(self, hidden_size1, hidden_size2, noise, alpha, step, max_epoch, lr
         self.step = step
         self.max_epoch = max_epoch
         self.lr = lr
-        self.device = "cuda" if torch.cuda.is_available() and not cpu else "cpu"
+        self.cpu = cpu
 
     def scale_matrix(self, mat):
         mat = mat - np.diag(np.diag(mat))
@@ -124,9 +124,10 @@ def get_emb(self, matrix):
         return emb_matrix
 
     def train(self, graph, return_dict=False):
-        return self.forward(graph, return_dict=False)
+        return self.forward(graph, return_dict=return_dict)
 
     def forward(self, graph, return_dict=False):
+        device = "cuda" if torch.cuda.is_available() and not self.cpu else "cpu"
         G = graph.to_networkx()
         self.num_node = G.number_of_nodes()
         A = nx.adjacency_matrix(G).todense()
@@ -137,8 +138,8 @@ def forward(self, graph, return_dict=False):
         input_mat = torch.from_numpy(self.get_denoised_matrix(PPMI).astype(np.float32))
         model = DNGR_layer(self.num_node, self.hidden_size1, self.hidden_size2)
 
-        input_mat = input_mat.to(self.device)
-        model = model.to(self.device)
+        input_mat = input_mat.to(device)
+        model = model.to(device)
 
         opt = torch.optim.Adam(model.parameters(), lr=self.lr)
         loss_func = nn.MSELoss()
diff --git a/cogdl/models/emb/gatne.py b/cogdl/models/emb/gatne.py
index ad3e666f..74bc084b 100644
--- a/cogdl/models/emb/gatne.py
+++ b/cogdl/models/emb/gatne.py
@@ -105,14 +105,13 @@ def __init__(
         self.num_sampled = negative_samples
         self.neighbor_samples = neighbor_samples
         self.schema = schema
-
         self.multiplicity = True
-        self.device = "cpu" if not torch.cuda.is_available() else "cuda"
 
     def train(self, network_data):
         return self.forward(network_data)
 
     def forward(self, network_data):
+        device = "cpu" if not torch.cuda.is_available() else "cuda"
         all_walks = generate_walks(network_data, self.walk_num, self.walk_length, schema=self.schema)
         vocab, index2word = generate_vocab(all_walks)
         train_pairs = generate_pairs(all_walks, vocab)
@@ -156,8 +155,8 @@ def forward(self, network_data):
         model = GATNEModel(num_nodes, embedding_size, embedding_u_size, edge_type_count, dim_att)
         nsloss = NSLoss(num_nodes, num_sampled, embedding_size)
 
-        model.to(self.device)
-        nsloss.to(self.device)
+        model.to(device)
+        nsloss.to(device)
 
         optimizer = torch.optim.Adam([{"params": model.parameters()}, {"params": nsloss.parameters()}], lr=1e-4)
 
@@ -176,11 +175,11 @@ def forward(self, network_data):
             for i, data in enumerate(data_iter):
                 optimizer.zero_grad()
                 embs = model(
-                    data[0].to(self.device),
-                    data[2].to(self.device),
-                    data[3].to(self.device),
+                    data[0].to(device),
+                    data[2].to(device),
+                    data[3].to(device),
                 )
-                loss = nsloss(data[0].to(self.device), embs, data[1].to(self.device))
+                loss = nsloss(data[0].to(device), embs, data[1].to(device))
                 loss.backward()
                 optimizer.step()
 
@@ -197,9 +196,9 @@ def forward(self, network_data):
 
         final_model = dict(zip(edge_types, [dict() for _ in range(edge_type_count)]))
         for i in range(num_nodes):
-            train_inputs = torch.tensor([i for _ in range(edge_type_count)]).to(self.device)
-            train_types = torch.tensor(list(range(edge_type_count))).to(self.device)
-            node_neigh = torch.tensor([neighbors[i] for _ in range(edge_type_count)]).to(self.device)
+            train_inputs = torch.tensor([i for _ in range(edge_type_count)]).to(device)
+            train_types = torch.tensor(list(range(edge_type_count))).to(device)
+            node_neigh = torch.tensor([neighbors[i] for _ in range(edge_type_count)]).to(device)
             node_emb = model(train_inputs, train_types, node_neigh)
             for j in range(edge_type_count):
                 final_model[edge_types[j]][index2word[i]] = node_emb[j].cpu().detach().numpy()
diff --git a/cogdl/models/emb/hin2vec.py b/cogdl/models/emb/hin2vec.py
index 0c1ad02e..9782674f 100644
--- a/cogdl/models/emb/hin2vec.py
+++ b/cogdl/models/emb/hin2vec.py
@@ -171,13 +171,13 @@ def __init__(self, hidden_dim, walk_length, walk_num, batch_size, hop, negative,
         self.negative = negative
         self.epochs = epochs
         self.lr = lr
-
-        self.device = "cpu" if not torch.cuda.is_available() or cpu else "cuda"
+        self.cpu = cpu
 
     def forward(self, data):
         return self.train(data)
 
     def train(self, data):
+        device = "cpu" if not torch.cuda.is_available() or self.cpu else "cuda"
         G = nx.DiGraph()
         row, col = data.edge_index
         G.add_edges_from(list(zip(row.numpy(), col.numpy())))
@@ -187,8 +187,8 @@ def train(self, data):
         pairs, relation = rw.data_preparation(walks, self.hop, self.negative)
 
         self.num_relation = len(relation)
-        model = Hin2vec_layer(self.num_node, self.num_relation, self.hidden_dim, self.device)
-        self.model = model.to(self.device)
+        model = Hin2vec_layer(self.num_node, self.num_relation, self.hidden_dim, device)
+        self.model = model.to(device)
 
         num_batch = int(len(pairs) / self.batch_size)
         print_num_batch = 100
@@ -200,7 +200,7 @@ def train(self, data):
             loss_n, pred, label = [], [], []
             for i in range(num_batch):
                 batch_pairs = torch.from_numpy(pairs[i * self.batch_size : (i + 1) * self.batch_size])
-                batch_pairs = batch_pairs.to(self.device)
+                batch_pairs = batch_pairs.to(device)
                 batch_pairs = batch_pairs.T
                 x, y, r, l = batch_pairs[0], batch_pairs[1], batch_pairs[2], batch_pairs[3]  # noqa E741
                 opt.zero_grad()
@@ -210,7 +210,7 @@ def train(self, data):
                 label.append(l)
                 pred.extend(logits)
                 if i % print_num_batch == 0 and i != 0:
-                    label = torch.cat(label).to(self.device)
+                    label = torch.cat(label).to(device)
                     pred = torch.stack(pred, dim=0)
                     pred = pred.max(1)[1]
                     acc = pred.eq(label).sum().item() / len(label)
diff --git a/cogdl/models/emb/line.py b/cogdl/models/emb/line.py
index 91aa4918..db2c706f 100644
--- a/cogdl/models/emb/line.py
+++ b/cogdl/models/emb/line.py
@@ -66,7 +66,7 @@ def __init__(self, dimension, walk_length, walk_num, negative, batch_size, alpha
         self.order = order
 
     def train(self, graph, return_dict=False):
-        return self.train(graph, return_dict)
+        return self.forward(graph, return_dict)
 
     def forward(self, graph, return_dict=False):
         # run LINE algorithm, 1-order, 2-order or 3(1-order + 2-order)
diff --git a/cogdl/models/emb/sdne.py b/cogdl/models/emb/sdne.py
index a2061089..8c89e9a8 100644
--- a/cogdl/models/emb/sdne.py
+++ b/cogdl/models/emb/sdne.py
@@ -105,12 +105,13 @@ def __init__(self, hidden_size1, hidden_size2, droput, alpha, beta, nu1, nu2, ma
         self.nu2 = nu2
         self.max_epoch = max_epoch
         self.lr = lr
-        self.device = "cuda" if torch.cuda.is_available() and not cpu else "cpu"
+        self.cpu = cpu
 
     def train(self, graph, return_dict=False):
         return self.forward(graph, return_dict)
 
     def forward(self, graph, return_dict=False):
+        device = "cuda" if torch.cuda.is_available() and not self.cpu else "cpu"
         G = graph.to_networkx()
         num_node = G.number_of_nodes()
         model = SDNE_layer(
@@ -120,8 +121,8 @@ def forward(self, graph, return_dict=False):
         A = torch.from_numpy(nx.adjacency_matrix(G).todense().astype(np.float32))
         L = torch.from_numpy(nx.laplacian_matrix(G).todense().astype(np.float32))
 
-        A, L = A.to(self.device), L.to(self.device)
-        model = model.to(self.device)
+        A, L = A.to(device), L.to(device)
+        model = model.to(device)
 
         opt = torch.optim.Adam(model.parameters(), lr=self.lr)
         epoch_iter = tqdm(range(self.max_epoch))
diff --git a/cogdl/models/nn/agc.py b/cogdl/models/nn/agc.py
index f5801610..a4604e16 100644
--- a/cogdl/models/nn/agc.py
+++ b/cogdl/models/nn/agc.py
@@ -31,11 +31,11 @@ def __init__(self, num_clusters, max_iter, cpu):
 
         self.num_clusters = num_clusters
         self.max_iter = max_iter
-
-        self.device = "cuda" if torch.cuda.is_available() and not cpu else "cpu"
+        self.cpu = cpu
 
     def forward(self, data):
-        data = data.to(self.device)
+        device = "cuda" if torch.cuda.is_available() and not self.cpu else "cpu"
+        data = data.to(device)
         self.num_nodes = data.x.shape[0]
         graph = data
         graph.add_remaining_self_loops()
diff --git a/cogdl/models/nn/mvgrl.py b/cogdl/models/nn/mvgrl.py
index f35f473c..bac796a8 100644
--- a/cogdl/models/nn/mvgrl.py
+++ b/cogdl/models/nn/mvgrl.py
@@ -157,7 +157,6 @@ def preprocess(self, graph):
 
         self.cache["diff"] = graphs[1]
         self.cache["adj"] = graphs[0]
-        self.device = next(self.gcn1.parameters()).device
         print("Preprocessing Done...")
 
     def forward(self, graph):
@@ -169,6 +168,7 @@ def forward(self, graph):
             self.preprocess(graph)
         diff, adj = self.cache["diff"], self.cache["adj"]
 
+        self.sample_size = min(self.sample_size, graph.num_nodes - self.batch_size)
         idx = np.random.randint(0, graph.num_nodes - self.sample_size + 1, self.batch_size)
         logits = []
         for i in idx:
diff --git a/cogdl/models/nn/unsup_graphsage.py b/cogdl/models/nn/unsup_graphsage.py
index c0d3662f..be796f1f 100644
--- a/cogdl/models/nn/unsup_graphsage.py
+++ b/cogdl/models/nn/unsup_graphsage.py
@@ -38,11 +38,6 @@ def add_args(parser):
         parser.add_argument("--num-layers", type=int, default=2)
         parser.add_argument("--sample-size", type=int, nargs='+', default=[10, 10])
         parser.add_argument("--dropout", type=float, default=0.5)
-        parser.add_argument("--walk-length", type=int, default=10)
-        parser.add_argument("--negative-samples", type=int, default=30)
-        parser.add_argument("--lr", type=float, default=0.001)
-
-        parser.add_argument("--max-epochs", type=int, default=3000)
         # fmt: on
 
     @classmethod
@@ -53,11 +48,9 @@ def build_model_from_args(cls, args):
             num_layers=args.num_layers,
             sample_size=args.sample_size,
             dropout=args.dropout,
-            walk_length=args.walk_length,
-            negative_samples=args.negative_samples,
         )
 
-    def __init__(self, num_features, hidden_size, num_layers, sample_size, dropout, walk_length, negative_samples):
+    def __init__(self, num_features, hidden_size, num_layers, sample_size, dropout):
         super(SAGE, self).__init__()
         self.adjlist = {}
         self.num_features = num_features
@@ -65,16 +58,9 @@ def __init__(self, num_features, hidden_size, num_layers, sample_size, dropout,
         self.num_layers = num_layers
         self.sample_size = sample_size
         self.dropout = dropout
-        self.walk_length = walk_length
-        self.num_negative_samples = negative_samples
-        self.walk_res = None
-        self.num_nodes = 0
-        self.negative_samples = 1
 
         shapes = [num_features] + [hidden_size] * num_layers
-
         self.convs = nn.ModuleList([SAGELayer(shapes[layer], shapes[layer + 1]) for layer in range(num_layers)])
-        self.random_walker = RandomWalker()
 
     def forward(self, graph):
         x = graph.x
@@ -88,34 +74,6 @@ def forward(self, graph):
                     x = F.dropout(x, p=self.dropout, training=self.training)
         return x
 
-    def loss(self, data):
-        x = self.forward(data)
-        device = x.device
-
-        self.random_walker.build_up(data.edge_index, data.x.shape[0])
-        walk_res = self.random_walker.walk(
-            start=torch.arange(0, x.shape[0]).to(device), walk_length=self.walk_length + 1
-        )
-        self.walk_res = torch.as_tensor(walk_res)[:, 1:]
-
-        if not self.num_nodes:
-            self.num_nodes = max(data.edge_index[0].max(), data.edge_index[1].max()).item() + 1
-
-        # if self.negative_samples is None:
-        self.negative_samples = torch.from_numpy(
-            np.random.choice(self.num_nodes, (self.num_nodes, self.num_negative_samples))
-        ).to(device)
-
-        pos_loss = -torch.log(
-            torch.sigmoid(torch.sum(x.unsqueeze(1).repeat(1, self.walk_length, 1) * x[self.walk_res], dim=-1))
-        ).mean()
-        neg_loss = -torch.log(
-            torch.sigmoid(
-                -torch.sum(x.unsqueeze(1).repeat(1, self.num_negative_samples, 1) * x[self.negative_samples], dim=-1)
-            )
-        ).mean()
-        return (pos_loss + neg_loss) / 2
-
     def embed(self, data):
         emb = self.forward(data)
         return emb
diff --git a/cogdl/pipelines.py b/cogdl/pipelines.py
index 15064b8c..17829896 100644
--- a/cogdl/pipelines.py
+++ b/cogdl/pipelines.py
@@ -189,7 +189,7 @@ def __call__(self, edge_index, x=None, edge_weight=None):
             num_nodes = edge_index.max().item() + 1
             if x is None:
                 print("No input node features, use random features instead.")
-                np.random.randn(num_nodes, self.num_features)
+                x = np.random.randn(num_nodes, self.num_features)
             if isinstance(x, np.ndarray):
                 x = torch.from_numpy(x).float()
             if isinstance(edge_index, np.ndarray):
@@ -200,7 +200,8 @@ def __call__(self, edge_index, x=None, edge_weight=None):
             dataset = NodeDataset(path=self.data_path, scale_feat=False, metric="accuracy")
             self.args.dataset = dataset
             model = train(self.args)
-            embeddings = model.embed(data).cpu().numpy()
+            embeddings = model.embed(data.to(model.device))
+            embeddings = embeddings.detach().cpu().numpy()
 
         return embeddings
 
diff --git a/cogdl/wrappers/default_match.py b/cogdl/wrappers/default_match.py
index b7d6358a..799358ca 100644
--- a/cogdl/wrappers/default_match.py
+++ b/cogdl/wrappers/default_match.py
@@ -45,7 +45,7 @@ def set_default_wrapper_config():
     network_embedding_models = [
         "deepwalk",
         "line",
-        "node2veec",
+        "node2vec",
         "prone",
         "netmf",
         "netsmf",
@@ -77,6 +77,7 @@ def set_default_wrapper_config():
     node_classification_wrappers["dgi"]["mw"] = "dgi_mw"
     node_classification_wrappers["m3s"]["mw"] = "m3s_mw"
     node_classification_wrappers["graphsage"]["mw"] = "graphsage_mw"
+    node_classification_wrappers["unsup_graphsage"]["mw"] = "unsup_graphsage_mw"
     node_classification_wrappers["mvgrl"]["mw"] = "mvgrl_mw"
     node_classification_wrappers["sagn"]["mw"] = "sagn_mw"
     node_classification_wrappers["grand"]["mw"] = "grand_mw"
diff --git a/cogdl/wrappers/model_wrapper/__init__.py b/cogdl/wrappers/model_wrapper/__init__.py
index 408378b6..2727c80c 100644
--- a/cogdl/wrappers/model_wrapper/__init__.py
+++ b/cogdl/wrappers/model_wrapper/__init__.py
@@ -37,6 +37,7 @@ def fetch_model_wrapper(name):
     "mvgrl_mw": "cogdl.wrappers.model_wrapper.node_classification.MVGRLModelWrapper",
     "self_auxiliary_mw": "cogdl.wrappers.model_wrapper.node_classification.SelfAuxiliaryModelWrapper",
     "graphsage_mw": "cogdl.wrappers.model_wrapper.node_classification.GraphSAGEModelWrapper",
+    "unsup_graphsage_mw": "cogdl.wrappers.model_wrapper.node_classification.UnsupGraphSAGEModelWrapper",
     "m3s_mw": "cogdl.wrappers.model_wrapper.node_classification.M3SModelWrapper",
     "network_embedding_mw": "cogdl.wrappers.model_wrapper.node_classification.NetworkEmbeddingModelWrapper",
     "node_classification_mw": "cogdl.wrappers.model_wrapper.node_classification.NodeClfModelWrapper",
diff --git a/cogdl/wrappers/model_wrapper/node_classification/__init__.py b/cogdl/wrappers/model_wrapper/node_classification/__init__.py
index b28c97a3..005e6145 100644
--- a/cogdl/wrappers/model_wrapper/node_classification/__init__.py
+++ b/cogdl/wrappers/model_wrapper/node_classification/__init__.py
@@ -4,6 +4,7 @@
 from .grace_mw import GRACEModelWrapper
 from .grand_mw import GrandModelWrapper
 from .graphsage_mw import GraphSAGEModelWrapper
+from .unsup_graphsage_mw import UnsupGraphSAGEModelWrapper
 from .m3s_mw import M3SModelWrapper
 from .mvgrl_mw import MVGRLModelWrapper
 from .network_embedding_mw import NetworkEmbeddingModelWrapper
diff --git a/cogdl/wrappers/model_wrapper/node_classification/unsup_graphsage_mw.py b/cogdl/wrappers/model_wrapper/node_classification/unsup_graphsage_mw.py
new file mode 100644
index 00000000..5b1447fc
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/node_classification/unsup_graphsage_mw.py
@@ -0,0 +1,61 @@
+import torch
+
+import numpy as np
+from cogdl.utils import RandomWalker
+from cogdl.wrappers.tools.wrapper_utils import evaluate_node_embeddings_using_logreg
+from .. import ModelWrapper
+
+
+class UnsupGraphSAGEModelWrapper(ModelWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--walk-length", type=int, default=10)
+        parser.add_argument("--negative-samples", type=int, default=30)
+        # fmt: on
+
+    def __init__(self, model, optimizer_cfg, walk_length, negative_samples):
+        super(UnsupGraphSAGEModelWrapper, self).__init__()
+        self.model = model
+        self.optimizer_cfg = optimizer_cfg
+        self.walk_length = walk_length
+        self.num_negative_samples = negative_samples
+        self.random_walker = RandomWalker()
+
+    def train_step(self, batch):
+        data = batch
+        x = self.model(data)
+        device = x.device
+
+        self.random_walker.build_up(data.edge_index, data.x.shape[0])
+        walk_res = self.random_walker.walk(
+            start=torch.arange(0, x.shape[0]).to(device), walk_length=self.walk_length + 1
+        )
+        self.walk_res = torch.as_tensor(walk_res)[:, 1:]
+
+        self.num_nodes = max(data.edge_index[0].max(), data.edge_index[1].max()).item() + 1
+
+        self.negative_samples = torch.from_numpy(
+            np.random.choice(self.num_nodes, (self.num_nodes, self.num_negative_samples))
+        ).to(device)
+
+        pos_loss = -torch.log(
+            torch.sigmoid(torch.sum(x.unsqueeze(1).repeat(1, self.walk_length, 1) * x[self.walk_res], dim=-1))
+        ).mean()
+        neg_loss = -torch.log(
+            torch.sigmoid(
+                -torch.sum(x.unsqueeze(1).repeat(1, self.num_negative_samples, 1) * x[self.negative_samples], dim=-1)
+            )
+        ).mean()
+        return (pos_loss + neg_loss) / 2
+
+    def test_step(self, graph):
+        with torch.no_grad():
+            pred = self.model(graph)
+        y = graph.y
+        result = evaluate_node_embeddings_using_logreg(pred, y, graph.train_mask, graph.test_mask)
+        self.note("test_acc", result)
+
+    def setup_optimizer(self):
+        cfg = self.optimizer_cfg
+        return torch.optim.Adam(self.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"])
diff --git a/docs/source/api/model_wrappers.rst b/docs/source/api/model_wrappers.rst
index b0a38b23..2f05ea76 100644
--- a/docs/source/api/model_wrappers.rst
+++ b/docs/source/api/model_wrappers.rst
@@ -39,6 +39,11 @@ Node Classification
     :undoc-members:
     :show-inheritance:
 
+.. autoclass:: cogdl.wrappers.model_wrapper.node_classification.UnsupGraphSAGEModelWrapper
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
 .. autoclass:: cogdl.wrappers.model_wrapper.node_classification.M3SModelWrapper
     :members:
     :undoc-members:
diff --git a/examples/generate_emb.py b/examples/generate_emb.py
index 6d6210c3..0c552636 100644
--- a/examples/generate_emb.py
+++ b/examples/generate_emb.py
@@ -17,6 +17,6 @@
 
 # build a pipeline for generating embeddings using unsupervised GNNs
 # pass model name and num_features with its hyper-parameters to this API
-generator = pipeline("generate-emb", model="mvgrl", no_test=True, num_features=8, hidden_size=4, sample_size=2)
+generator = pipeline("generate-emb", model="mvgrl", no_test=True, num_features=8, hidden_size=4)
 outputs = generator(edge_index, x=np.random.randn(8, 8))
 print(outputs)


Test patch:
diff --git a/tests/models/ssl/test_contrastive_models.py b/tests/models/ssl/test_contrastive_models.py
index 58a8fdce..1264dc96 100644
--- a/tests/models/ssl/test_contrastive_models.py
+++ b/tests/models/ssl/test_contrastive_models.py
@@ -42,7 +42,7 @@ def get_default_args_for_unsup_nn(dataset, model, dw="node_classification_dw", m
 
 
 def test_unsupervised_graphsage():
-    args = get_default_args_for_unsup_nn("cora", "unsup_graphsage")
+    args = get_default_args_for_unsup_nn("cora", "unsup_graphsage", mw="unsup_graphsage_mw")
     args.negative_samples = 10
     args.walk_length = 5
     args.sample_size = [5, 5]
@@ -91,6 +91,6 @@ def test_grace():
 
 if __name__ == "__main__":
     test_unsupervised_graphsage()
-    test_grace()
-    test_mvgrl()
-    test_dgi()
+    # test_grace()
+    # test_mvgrl()
+    # test_dgi()