fix some bugs in link_prediction

aimclub · Dec 19, 2023 · b7317da · b7317da
1 parent cb0a70c
commit b7317da
Show file tree

Hide file tree

Showing 6 changed files with 2,343 additions and 147 deletions.
diff --git a/stable_gnn/model_gc.py b/stable_gnn/model_gc.py
@@ -84,6 +84,7 @@ def forward(self, x: Tensor, edge_index: Adj, batch: Tensor) -> Tuple[Tensor, Te
         :return: (Tensor, Tensor): Predicted probabilities of labels and predicted degrees of graphs
         """
         # 1. Obtain node embeddings
+        x = x.type(torch.FloatTensor).to(self.device)
         for i, conv in enumerate(self.convs):
             x = conv(x, edge_index)
             if i < self.num_layers - 1:
@@ -111,7 +112,7 @@ def loss_sup(pred: Tensor, label: Tensor) -> Tensor:
         :param label: (Tensor): Genuine labels
         :return: (Tensor): Loss
         """
-        return F.nll_loss(pred, label)
+        return F.nll_loss(pred, label)#.type(torch.LongTensor).to(device))
 
     @staticmethod
     def convert_dataset(

diff --git a/stable_gnn/model_link_predict.py b/stable_gnn/model_link_predict.py
@@ -78,7 +78,7 @@ def train_cl(self, train_edges: List[List[int]], neg_samples_train: List[List[in
         :param neg_samples_train: (List): List of negative samples to train
         :return: (BaseEstimator): Classifier which support fit predict notation
         """
-        if self.number_of_trials:
+        if self.number_of_trials > 0:
             self.embeddings = EmbeddingFactory().build_embeddings(
                 loss_name=self.loss_name,
                 conv=self.emb_conv_name,