Merge pull request #17 from florencejt/bugs/attentionGNNmulticlass

Bug fix in AttentionWeightedGNN with multiclass classification

fusilli/fusionmodels/tabularfusion/attention_weighted_GNN.py

@@ -163,6 +163,10 @@ def training_step(self, batch, batch_idx):
 
         y_hat, weights = self.forward((x1, x2))
 
+        if self.prediction_task == "multiclass":
+            # turn the labels into one hot vectors
+            y = F.one_hot(y, num_classes=self.multiclass_dimensions).to(torch.float32)
+
         loss = F.mse_loss(y_hat.squeeze(), y.to(torch.float32).squeeze())
         self.log('train_loss', loss, logger=None)
         return loss
@@ -186,6 +190,11 @@ def validation_step(self, batch, batch_idx):
         """
         x1, x2, y = batch
         y_hat, weights = self.forward((x1, x2))
+
+        if self.prediction_task == "multiclass":
+            # turn the labels into one hot vectors
+            y = F.one_hot(y, num_classes=self.multiclass_dimensions).to(torch.float32)
+
         loss = F.mse_loss(y_hat.squeeze(), y.to(torch.float32).squeeze())
         self.log('val_loss', loss, logger=None)
         return loss