fix size_nn

app.py

@@ -72,7 +72,7 @@ print('Loaded diffusion model')
 
 
 def sample_fn(_data):
-    output, _ = size_nn.forward(_data)
+    output, _ = size_nn.forward(_data, return_loss=False)
     probabilities = torch.softmax(output, dim=1)
     distribution = torch.distributions.Categorical(probs=probabilities)
     samples = distribution.sample()

src/linker_size_lightning.py

@@ -79,7 +79,7 @@ class SizeClassifier(pl.LightningModule):
     def test_dataloader(self):
         return get_dataloader(self.test_dataset, self.batch_size, collate_fn=collate_with_fragment_edges)
 
-    def forward(self, data):
+    def forward(self, data, return_loss=True):
         h = data['one_hot']
         x = data['positions']
         fragment_mask = data['fragment_mask']
@@ -103,8 +103,11 @@ class SizeClassifier(pl.LightningModule):
         output = self.gnn.forward(h, edges, distances, fragment_mask, distance_edge_mask)
         output = output.view(bs, n_nodes, -1).mean(1)
 
-        true = self.get_true_labels(linker_mask)
-        loss = cross_entropy(output, true, weight=self.loss_weights)
+        if return_loss:
+            true = self.get_true_labels(linker_mask)
+            loss = cross_entropy(output, true, weight=self.loss_weights)
+        else:
+            loss = None
 
         return output, loss